作者: Sayak Paul
创建日期 2021/04/13
上次修改日期 2021/04/19
描述:使用一致性正则化进行训练,以增强模型对数据分布变化的鲁棒性。
当数据独立同分布(i.i.d.)时,深度学习模型在许多图像识别任务中表现出色。但是,它们可能会受到输入数据中细微分布变化(例如随机噪声、对比度变化和模糊)引起的性能下降的影响。因此,自然而然地产生了一个问题,即为什么会出现这种情况。正如在计算机视觉中模型鲁棒性的傅里叶视角中所讨论的那样,深度学习模型没有理由对这种变化具有鲁棒性。标准的模型训练过程(例如标准的图像分类训练流程)无法使模型学习超出训练数据形式提供的范围。
在本示例中,我们将通过执行以下操作来训练一个图像分类模型,并在其中强制执行一种一致性:
此整体训练工作流程的根源在于诸如FixMatch、用于一致性训练的无监督数据增强和噪声学生训练等工作。由于此训练过程鼓励模型对干净和噪声图像产生一致的预测,因此它通常被称为一致性训练或使用一致性正则化进行训练。尽管此示例侧重于使用一致性训练来增强模型对常见损坏的鲁棒性,但此示例也可以用作执行弱监督学习的模板。
此示例需要 TensorFlow 2.4 或更高版本,以及 TensorFlow Hub 和 TensorFlow Models,可以使用以下命令安装它们:
!pip install -q tf-models-official tensorflow-addons
from official.vision.image_classification.augment import RandAugment
from tensorflow.keras import layers
import tensorflow as tf
import tensorflow_addons as tfa
import matplotlib.pyplot as plt
tf.random.set_seed(42)
AUTO = tf.data.AUTOTUNE
BATCH_SIZE = 128
EPOCHS = 5
CROP_TO = 72
RESIZE_TO = 96
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
val_samples = 49500
new_train_x, new_y_train = x_train[: val_samples + 1], y_train[: val_samples + 1]
val_x, val_y = x_train[val_samples:], y_train[val_samples:]
Dataset
对象# Initialize `RandAugment` object with 2 layers of
# augmentation transforms and strength of 9.
augmenter = RandAugment(num_layers=2, magnitude=9)
为了训练教师模型,我们将只使用两种几何增强变换:随机水平翻转和随机裁剪。
def preprocess_train(image, label, noisy=True):
image = tf.image.random_flip_left_right(image)
# We first resize the original image to a larger dimension
# and then we take random crops from it.
image = tf.image.resize(image, [RESIZE_TO, RESIZE_TO])
image = tf.image.random_crop(image, [CROP_TO, CROP_TO, 3])
if noisy:
image = augmenter.distort(image)
return image, label
def preprocess_test(image, label):
image = tf.image.resize(image, [CROP_TO, CROP_TO])
return image, label
train_ds = tf.data.Dataset.from_tensor_slices((new_train_x, new_y_train))
validation_ds = tf.data.Dataset.from_tensor_slices((val_x, val_y))
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
我们确保train_clean_ds
和train_noisy_ds
使用相同的种子进行混洗,以确保它们的顺序完全相同。这在训练学生模型期间将非常有用。
# This dataset will be used to train the first model.
train_clean_ds = (
train_ds.shuffle(BATCH_SIZE * 10, seed=42)
.map(lambda x, y: (preprocess_train(x, y, noisy=False)), num_parallel_calls=AUTO)
.batch(BATCH_SIZE)
.prefetch(AUTO)
)
# This prepares the `Dataset` object to use RandAugment.
train_noisy_ds = (
train_ds.shuffle(BATCH_SIZE * 10, seed=42)
.map(preprocess_train, num_parallel_calls=AUTO)
.batch(BATCH_SIZE)
.prefetch(AUTO)
)
validation_ds = (
validation_ds.map(preprocess_test, num_parallel_calls=AUTO)
.batch(BATCH_SIZE)
.prefetch(AUTO)
)
test_ds = (
test_ds.map(preprocess_test, num_parallel_calls=AUTO)
.batch(BATCH_SIZE)
.prefetch(AUTO)
)
# This dataset will be used to train the second model.
consistency_training_ds = tf.data.Dataset.zip((train_clean_ds, train_noisy_ds))
sample_images, sample_labels = next(iter(train_clean_ds))
plt.figure(figsize=(10, 10))
for i, image in enumerate(sample_images[:9]):
ax = plt.subplot(3, 3, i + 1)
plt.imshow(image.numpy().astype("int"))
plt.axis("off")
sample_images, sample_labels = next(iter(train_noisy_ds))
plt.figure(figsize=(10, 10))
for i, image in enumerate(sample_images[:9]):
ax = plt.subplot(3, 3, i + 1)
plt.imshow(image.numpy().astype("int"))
plt.axis("off")
现在我们定义我们的模型构建实用程序。我们的模型基于ResNet50V2 架构。
def get_training_model(num_classes=10):
resnet50_v2 = tf.keras.applications.ResNet50V2(
weights=None, include_top=False, input_shape=(CROP_TO, CROP_TO, 3),
)
model = tf.keras.Sequential(
[
layers.Input((CROP_TO, CROP_TO, 3)),
layers.Rescaling(scale=1.0 / 127.5, offset=-1),
resnet50_v2,
layers.GlobalAveragePooling2D(),
layers.Dense(num_classes),
]
)
return model
为了可重复性,我们序列化教师网络的初始随机权重。
initial_teacher_model = get_training_model()
initial_teacher_model.save_weights("initial_teacher_model.h5")
如“噪声学生训练”中所述,如果教师模型使用几何集成进行训练,并且当学生模型被迫模仿它时,会导致更好的性能。最初的工作使用随机深度和Dropout来引入集成部分,但在这个例子中,我们将使用随机权重平均 (SWA),它也类似于几何集成。
# Define the callbacks.
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(patience=3)
early_stopping = tf.keras.callbacks.EarlyStopping(
patience=10, restore_best_weights=True
)
# Initialize SWA from tf-hub.
SWA = tfa.optimizers.SWA
# Compile and train the teacher model.
teacher_model = get_training_model()
teacher_model.load_weights("initial_teacher_model.h5")
teacher_model.compile(
# Notice that we are wrapping our optimizer within SWA
optimizer=SWA(tf.keras.optimizers.Adam()),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=["accuracy"],
)
history = teacher_model.fit(
train_clean_ds,
epochs=EPOCHS,
validation_data=validation_ds,
callbacks=[reduce_lr, early_stopping],
)
# Evaluate the teacher model on the test set.
_, acc = teacher_model.evaluate(test_ds, verbose=0)
print(f"Test accuracy: {acc*100}%")
Epoch 1/5
387/387 [==============================] - 73s 78ms/step - loss: 1.7785 - accuracy: 0.3582 - val_loss: 2.0589 - val_accuracy: 0.3920
Epoch 2/5
387/387 [==============================] - 28s 71ms/step - loss: 1.2493 - accuracy: 0.5542 - val_loss: 1.4228 - val_accuracy: 0.5380
Epoch 3/5
387/387 [==============================] - 28s 73ms/step - loss: 1.0294 - accuracy: 0.6350 - val_loss: 1.4422 - val_accuracy: 0.5900
Epoch 4/5
387/387 [==============================] - 28s 73ms/step - loss: 0.8954 - accuracy: 0.6864 - val_loss: 1.2189 - val_accuracy: 0.6520
Epoch 5/5
387/387 [==============================] - 28s 73ms/step - loss: 0.7879 - accuracy: 0.7231 - val_loss: 0.9790 - val_accuracy: 0.6500
Test accuracy: 65.83999991416931%
对于这部分,我们将借用这个 Keras 示例中的Distiller
类。
# Majority of the code is taken from:
# https://keras.org.cn/examples/vision/knowledge_distillation/
class SelfTrainer(tf.keras.Model):
def __init__(self, student, teacher):
super().__init__()
self.student = student
self.teacher = teacher
def compile(
self, optimizer, metrics, student_loss_fn, distillation_loss_fn, temperature=3,
):
super().compile(optimizer=optimizer, metrics=metrics)
self.student_loss_fn = student_loss_fn
self.distillation_loss_fn = distillation_loss_fn
self.temperature = temperature
def train_step(self, data):
# Since our dataset is a zip of two independent datasets,
# after initially parsing them, we segregate the
# respective images and labels next.
clean_ds, noisy_ds = data
clean_images, _ = clean_ds
noisy_images, y = noisy_ds
# Forward pass of teacher
teacher_predictions = self.teacher(clean_images, training=False)
with tf.GradientTape() as tape:
# Forward pass of student
student_predictions = self.student(noisy_images, training=True)
# Compute losses
student_loss = self.student_loss_fn(y, student_predictions)
distillation_loss = self.distillation_loss_fn(
tf.nn.softmax(teacher_predictions / self.temperature, axis=1),
tf.nn.softmax(student_predictions / self.temperature, axis=1),
)
total_loss = (student_loss + distillation_loss) / 2
# Compute gradients
trainable_vars = self.student.trainable_variables
gradients = tape.gradient(total_loss, trainable_vars)
# Update weights
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
# Update the metrics configured in `compile()`
self.compiled_metrics.update_state(
y, tf.nn.softmax(student_predictions, axis=1)
)
# Return a dict of performance
results = {m.name: m.result() for m in self.metrics}
results.update({"total_loss": total_loss})
return results
def test_step(self, data):
# During inference, we only pass a dataset consisting images and labels.
x, y = data
# Compute predictions
y_prediction = self.student(x, training=False)
# Update the metrics
self.compiled_metrics.update_state(y, tf.nn.softmax(y_prediction, axis=1))
# Return a dict of performance
results = {m.name: m.result() for m in self.metrics}
return results
此实现中唯一的区别是损失的计算方式。**我们不是以不同的方式对蒸馏损失和学生损失进行加权,而是按照噪声学生训练的方法取它们的平均值**。
# Define the callbacks.
# We are using a larger decay factor to stabilize the training.
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
patience=3, factor=0.5, monitor="val_accuracy"
)
early_stopping = tf.keras.callbacks.EarlyStopping(
patience=10, restore_best_weights=True, monitor="val_accuracy"
)
# Compile and train the student model.
self_trainer = SelfTrainer(student=get_training_model(), teacher=teacher_model)
self_trainer.compile(
# Notice we are *not* using SWA here.
optimizer="adam",
metrics=["accuracy"],
student_loss_fn=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
distillation_loss_fn=tf.keras.losses.KLDivergence(),
temperature=10,
)
history = self_trainer.fit(
consistency_training_ds,
epochs=EPOCHS,
validation_data=validation_ds,
callbacks=[reduce_lr, early_stopping],
)
# Evaluate the student model.
acc = self_trainer.evaluate(test_ds, verbose=0)
print(f"Test accuracy from student model: {acc*100}%")
Epoch 1/5
387/387 [==============================] - 39s 84ms/step - accuracy: 0.2112 - total_loss: 1.0629 - val_accuracy: 0.4180
Epoch 2/5
387/387 [==============================] - 32s 82ms/step - accuracy: 0.3341 - total_loss: 0.9554 - val_accuracy: 0.3900
Epoch 3/5
387/387 [==============================] - 31s 81ms/step - accuracy: 0.3873 - total_loss: 0.8852 - val_accuracy: 0.4580
Epoch 4/5
387/387 [==============================] - 31s 81ms/step - accuracy: 0.4294 - total_loss: 0.8423 - val_accuracy: 0.5660
Epoch 5/5
387/387 [==============================] - 31s 81ms/step - accuracy: 0.4547 - total_loss: 0.8093 - val_accuracy: 0.5880
Test accuracy from student model: 58.490002155303955%
评估视觉模型鲁棒性的一个标准基准是记录它们在损坏的数据集(如 ImageNet-C 和 CIFAR-10-C)上的性能,这两个数据集都在对常见损坏和扰动的神经网络鲁棒性进行基准测试中提出。对于此示例,我们将使用 CIFAR-10-C 数据集,该数据集对 5 个不同严重程度级别有 19 种不同的损坏。为了评估模型在此数据集上的鲁棒性,我们将执行以下操作
出于本示例的目的,我们不会执行这些步骤。这就是我们仅训练模型 5 个 epoch 的原因。您可以查看此存储库,该存储库演示了完整的训练实验以及前面提到的评估。下图显示了该评估的执行摘要
平均 Top-1 结果代表 CIFAR-10-C 数据集,而测试 Top-1 结果代表 CIFAR-10 测试集。很明显,一致性训练不仅在增强模型鲁棒性方面具有优势,而且在提高标准测试性能方面也具有优势。