自监督对比学习与 NNCLR
作者: Rishit Dagli
创建日期: 2021/09/13
最后修改: 2024/01/22
描述: NNCLR 的实现，一种用于计算机视觉的自监督学习方法。
介绍
自监督学习
自监督表示学习旨在从原始数据中获得样本的鲁棒表示，而无需昂贵的标签或注释。该领域早期的方法集中于定义预训练任务，这些任务涉及在拥有丰富弱监督标签的领域中的替代任务。经过训练以解决此类任务的编码器预计能够学习对其他需要昂贵注释的下游任务（如图像分类）有用的一般特征。
对比学习
一类广泛的自监督学习技术是使用 对比损失 的技术，这些技术已经应用于广泛的计算机视觉应用，如 图像相似性、降维 (DrLIM) 和 人脸验证/识别。这些方法学习一个潜在空间，该空间将正样本聚集在一起，同时将负样本推开。
NNCLR
在这个示例中，我们实现了 NNCLR，正如论文 With a Little Help from My Friends: Nearest-Neighbor Contrastive Learning of Visual Representations 中所提出的，来自 Google Research 和 DeepMind。
NNCLR 学习的自监督表示超越了单实例正样本，这允许学习对不同视角、变形甚至类内变异不变的更好特征。基于聚类的方法为超越单实例正样本提供了一个很好的方法，但假设整个聚类都是正样本可能会由于过早的过度概括而影响性能。相反，NNCLR 在学习的表示空间中使用最近邻作为正样本。此外，NNCLR 提高了现有对比学习方法的性能，如 SimCLR（Keras 示例），并减少了自监督方法对数据增强策略的依赖。
以下是论文作者提供的出色可视化，展示了 NNCLR 如何基于 SimCLR 的思想：
我们可以看到，SimCLR 使用同一图像的两个视图作为正样本对。这两个视图是使用随机数据增强生成的，经过编码器处理以获得正嵌入对，最终使用两个增强。NNCLR 则保持一个代表完整数据分布的 支持集 的嵌入，并使用最近邻形成正样本对。在训练期间，支持集被用作内存，类似于队列（即先进先出），如 MoCo 中所示。
此示例需要 tensorflow_datasets，可以使用以下命令安装：
!pip install tensorflow-datasets
Setup
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_datasets as tfds
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras
import keras_cv
from keras import ops
from keras import layers
超参数
更大的 queue_size 很可能意味着更好的性能，如原始论文所示，但会引入显著的计算开销。作者表明，NNCLR 的最佳结果是在队列大小为 98,304 时（他们实验过的最大 queue_size）。我们在这里使用 10,000 来展示一个工作示例。
AUTOTUNE = tf.data.AUTOTUNE
shuffle_buffer = 5000
# 以下两个值取自 https://www.tensorflow.org/datasets/catalog/stl10
labelled_train_images = 5000
unlabelled_images = 100000
temperature = 0.1
queue_size = 10000
contrastive_augmenter = {
"brightness": 0.5,
"name": "contrastive_augmenter",
"scale": (0.2, 1.0),
}
classification_augmenter = {
"brightness": 0.2,
"name": "classification_augmenter",
"scale": (0.5, 1.0),
}
input_shape = (96, 96, 3)
width = 128
num_epochs = 5 # 使用 25 可以获得更好的结果
steps_per_epoch = 50 # 使用 200 可以获得更好的结果
加载数据集
我们从 TensorFlow Datasets 加载 STL-10 数据集，这是一个用于开发无监督特征学习、深度学习、自学算法的图像识别数据集。它受到 CIFAR-10 数据集的启发，并进行了某些修改。
dataset_name = "stl10"
def prepare_dataset():
unlabeled_batch_size = unlabelled_images // steps_per_epoch
labeled_batch_size = labelled_train_images // steps_per_epoch
batch_size = unlabeled_batch_size + labeled_batch_size
unlabeled_train_dataset = (
tfds.load(
dataset_name, split="unlabelled", as_supervised=True, shuffle_files=True
)
.shuffle(buffer_size=shuffle_buffer)
.batch(unlabeled_batch_size, drop_remainder=True)
)
labeled_train_dataset = (
tfds.load(dataset_name, split="train", as_supervised=True, shuffle_files=True)
.shuffle(buffer_size=shuffle_buffer)
.batch(labeled_batch_size, drop_remainder=True)
)
test_dataset = (
tfds.load(dataset_name, split="test", as_supervised=True)
.batch(batch_size)
.prefetch(buffer_size=AUTOTUNE)
)
train_dataset = tf.data.Dataset.zip(
(unlabeled_train_dataset, labeled_train_dataset)
).prefetch(buffer_size=AUTOTUNE)
return batch_size, train_dataset, labeled_train_dataset, test_dataset
batch_size, train_dataset, labeled_train_dataset, test_dataset = prepare_dataset()
数据增强
其他自监督技术如 SimCLR， BYOL，SwAV 等， 在设计良好的数据增强管道上高度依赖以获得最佳性能。 然而，NNCLR 对复杂增强的依赖性较低，因为最近邻已经提供了样本变化的丰富性。一些常见的技术通常包括在增强管道中：
- 随机调整大小裁剪
- 多种颜色失真
- 高斯模糊
由于 NNCLR 对复杂增强的依赖性较低，我们将仅使用随机裁剪和随机亮度来增强输入图像。
准备增强模块
def augmenter(brightness, name, scale):
return keras.Sequential(
[
layers.Input(shape=input_shape),
layers.Rescaling(1 / 255),
layers.RandomFlip("horizontal"),
keras_cv.layers.RandomCropAndResize(
target_size=(input_shape[0], input_shape[1]),
crop_area_factor=scale,
aspect_ratio_factor=(3 / 4, 4 / 3),
),
keras_cv.layers.RandomBrightness(factor=brightness, value_range=(0.0, 1.0)),
],
name=name,
)
编码器架构
使用 ResNet-50 作为编码器架构是文献中的标准。在原始论文中，作者使用 ResNet-50 作为编码器架构，并对 ResNet-50 的输出进行空间平均。然而，请记住，更强大的模型不仅会增加训练时间，还会需要更多内存，并限制您可以使用的最大批量大小。为了这个示例的目的，我们仅使用四个卷积层。
def encoder():
return keras.Sequential(
[
layers.Input(shape=input_shape),
layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
layers.Conv2D(width, kernel_size=3, strides=2, activation="relu"),
layers.Flatten(),
layers.Dense(width, activation="relu"),
],
name="encoder",
)
NNCLR 模型用于对比预训练
我们在未标记的图像上训练一个编码器，使用对比损失。在编码器顶部附加一个非线性投影头，因为它提高了编码器表示的质量。
class NNCLR(keras.Model):
def __init__(
self, temperature, queue_size,
):
super().__init__()
self.probe_accuracy = keras.metrics.SparseCategoricalAccuracy()
self.correlation_accuracy = keras.metrics.SparseCategoricalAccuracy()
self.contrastive_accuracy = keras.metrics.SparseCategoricalAccuracy()
self.probe_loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
self.contrastive_augmenter = augmenter(**contrastive_augmenter)
self.classification_augmenter = augmenter(**classification_augmenter)
self.encoder = encoder()
self.projection_head = keras.Sequential(
[
layers.Input(shape=(width,)),
layers.Dense(width, activation="relu"),
layers.Dense(width),
],
name="projection_head",
)
self.linear_probe = keras.Sequential(
[layers.Input(shape=(width,)), layers.Dense(10)], name="linear_probe"
)
self.temperature = temperature
feature_dimensions = self.encoder.output_shape[1]
self.feature_queue = keras.Variable(
keras.utils.normalize(
keras.random.normal(shape=(queue_size, feature_dimensions)),
axis=1,
order=2,
),
trainable=False,
)
def compile(self, contrastive_optimizer, probe_optimizer, **kwargs):
super().compile(**kwargs)
self.contrastive_optimizer = contrastive_optimizer
self.probe_optimizer = probe_optimizer
def nearest_neighbour(self, projections):
support_similarities = ops.matmul(projections, ops.transpose(self.feature_queue))
nn_projections = ops.take(
self.feature_queue, ops.argmax(support_similarities, axis=1), axis=0
)
return projections + ops.stop_gradient(nn_projections - projections)
def update_contrastive_accuracy(self, features_1, features_2):
features_1 = keras.utils.normalize(features_1, axis=1, order=2)
features_2 = keras.utils.normalize(features_2, axis=1, order=2)
similarities = ops.matmul(features_1, ops.transpose(features_2))
batch_size = ops.shape(features_1)[0]
contrastive_labels = ops.arange(batch_size)
self.contrastive_accuracy.update_state(
ops.concatenate([contrastive_labels, contrastive_labels], axis=0),
ops.concatenate([similarities, ops.transpose(similarities)], axis=0),
)
def update_correlation_accuracy(self, features_1, features_2):
features_1 = (features_1 - ops.mean(features_1, axis=0)) / ops.std(
features_1, axis=0
)
features_2 = (features_2 - ops.mean(features_2, axis=0)) / ops.std(
features_2, axis=0
)
batch_size = ops.shape(features_1)[0]
cross_correlation = (
ops.matmul(ops.transpose(features_1), features_2) / batch_size
)
feature_dim = ops.shape(features_1)[1]
correlation_labels = ops.arange(feature_dim)
self.correlation_accuracy.update_state(
ops.concatenate([correlation_labels, correlation_labels], axis=0),
ops.concatenate(
[cross_correlation, ops.transpose(cross_correlation)], axis=0
),
)
def contrastive_loss(self, projections_1, projections_2):
projections_1 = keras.utils.normalize(projections_1, axis=1, order=2)
projections_2 = keras.utils.normalize(projections_2, axis=1, order=2)
similarities_1_2_1 = (
ops.matmul(
self.nearest_neighbour(projections_1), ops.transpose(projections_2)
)
/ self.temperature
)
similarities_1_2_2 = (
ops.matmul(
projections_2, ops.transpose(self.nearest_neighbour(projections_1))
)
/ self.temperature
)
similarities_2_1_1 = (
ops.matmul(
self.nearest_neighbour(projections_2), ops.transpose(projections_1)
)
/ self.temperature
)
similarities_2_1_2 = (
ops.matmul(
projections_1, ops.transpose(self.nearest_neighbour(projections_2))
)
/ self.temperature
)
batch_size = ops.shape(projections_1)[0]
contrastive_labels = ops.arange(batch_size)
loss = keras.losses.sparse_categorical_crossentropy(
ops.concatenate(
[
contrastive_labels,
contrastive_labels,
contrastive_labels,
contrastive_labels,
],
axis=0,
),
ops.concatenate(
[
similarities_1_2_1,
similarities_1_2_2,
similarities_2_1_1,
similarities_2_1_2,
],
axis=0,
),
from_logits=True,
)
self.feature_queue.assign(
ops.concatenate([projections_1, self.feature_queue[:-batch_size]], axis=0)
)
return loss
def train_step(self, data):
(unlabeled_images, _), (labeled_images, labels) = data
images = ops.concatenate((unlabeled_images, labeled_images), axis=0)
augmented_images_1 = self.contrastive_augmenter(images)
augmented_images_2 = self.contrastive_augmenter(images)
with tf.GradientTape() as tape:
features_1 = self.encoder(augmented_images_1)
features_2 = self.encoder(augmented_images_2)
projections_1 = self.projection_head(features_1)
projections_2 = self.projection_head(features_2)
contrastive_loss = self.contrastive_loss(projections_1, projections_2)
gradients = tape.gradient(
contrastive_loss,
self.encoder.trainable_weights + self.projection_head.trainable_weights,
)
self.contrastive_optimizer.apply_gradients(
zip(
gradients,
self.encoder.trainable_weights + self.projection_head.trainable_weights,
)
)
self.update_contrastive_accuracy(features_1, features_2)
self.update_correlation_accuracy(features_1, features_2)
preprocessed_images = self.classification_augmenter(labeled_images)
with tf.GradientTape() as tape:
features = self.encoder(preprocessed_images)
class_logits = self.linear_probe(features)
probe_loss = self.probe_loss(labels, class_logits)
gradients = tape.gradient(probe_loss, self.linear_probe.trainable_weights)
self.probe_optimizer.apply_gradients(
zip(gradients, self.linear_probe.trainable_weights)
)
self.probe_accuracy.update_state(labels, class_logits)
return {
"c_loss": contrastive_loss,
"c_acc": self.contrastive_accuracy.result(),
"r_acc": self.correlation_accuracy.result(),
"p_loss": probe_loss,
"p_acc": self.probe_accuracy.result(),
}
def test_step(self, data):
labeled_images, labels = data
preprocessed_images = self.classification_augmenter(
labeled_images, training=False
)
features = self.encoder(preprocessed_images, training=False)
class_logits = self.linear_probe(features, training=False)
probe_loss = self.probe_loss(labels, class_logits)
self.probe_accuracy.update_state(labels, class_logits)
return {"p_loss": probe_loss, "p_acc": self.probe_accuracy.result()}
预训练 NNCLR
我们使用论文中建议的 temperature 为 0.1 并且如前所述 queue_size 为 10,000 来训练网络。我们使用 Adam 作为我们的对比和探测优化器。在这个示例中，我们仅训练模型 30 个周期，但为了更好的性能，它应该训练更多的周期。
可以使用以下两种指标来监测预训练性能，我们也记录了它们（摘自 this Keras example）：
- 对比准确率：自监督指标，表示图像的表示与其不同增强版本的表示更相似的情况下的比例，而不是与当前批次中任何其他图像的表示相比。自监督指标可以用于超参数调优，即使没有标记示例的情况下。
- 线性探测准确率：线性探测是评估自监督分类器的常用指标。它是通过在编码器特征上训练的逻辑回归分类器的准确率计算得出的。在我们的例子中，这是通过在冻结的编码器上训练单个稠密层来实现的。请注意，与传统方法相反，其中分类器在预训练阶段之后进行训练，在此示例中我们在预训练期间训练它。这可能会稍微降低其准确性，但这样我们可以在训练过程中监控其值，这有助于实验和调试。
model = NNCLR(temperature=temperature, queue_size=queue_size)
model.compile(
contrastive_optimizer=keras.optimizers.Adam(),
probe_optimizer=keras.optimizers.Adam(),
jit_compile=False,
)
pretrain_history = model.fit(
train_dataset, epochs=num_epochs, validation_data=test_dataset
)
自监督学习在你仅能获得非常有限的标记训练数据但能够管理构建一个大规模的未标记数据集时特别有帮助，正如以往的方法如 SEER，SimCLR，SwAV 等所示。
你也应该查看这些论文的博客文章，它们清楚地显示，通过先在大型未标记数据集上进行预训练，然后在较小的标记数据集上进行微调，确实可以用少量类别标签取得良好结果：
你也被建议查看 原始论文。
非常感谢 Debidatta Dwibedi（谷歌研究），NNCLR 论文的主要作者，他对这个示例进行了极具洞察力的评审。这个示例也受到 SimCLR Keras 示例 的启发。