2024最新对比学习：从零构建高效模型的入门指南

1次阅读

共计 3990 个字符，预计需要花费 10 分钟才能阅读完成。

对比学习（Contrastive Learning）作为自监督学习的重要分支，在 2024 年已成为计算机视觉和自然语言处理领域的关键技术。它通过让模型学会区分相似和不相似的数据样本，无需人工标注即可学习到高质量的特征表示。这种方法在小样本学习（Few-shot Learning）场景下表现尤为突出——当标注数据稀缺时，利用大量无标注数据进行预训练可以显著提升模型性能。

在图像领域，对比学习已被成功应用于图像分类、目标检测等任务；在文本领域，它帮助提升了语义相似度计算、文本分类等任务的准确性。与监督学习相比，对比学习最大的优势在于能够利用海量无标注数据，降低对昂贵人工标注的依赖。

目前对比学习领域主要有几种主流框架：

SimCLR（Simple Contrastive Learning of Visual Representations）：谷歌提出的框架，结构简单但效果出色
MoCo（Momentum Contrast）：Facebook 提出的框架，使用动量编码器和内存库
BYOL（Bootstrap Your Own Latent）：不需要负样本的创新方法

对于初学者，建议从 SimCLR 入手。原因在于：

实现相对简单，不需要复杂的内存库或动量编码器
代码结构清晰，便于理解和修改
在标准 benchmark 上表现优秀
有大量开源实现和教程可供参考

数据增强是对比学习成功的关键。在 SimCLR 中，每张输入图像会经过两次不同的增强变换，生成一对正样本。常用的增强方法包括：

随机裁剪（Random Cropping）：在保持图像主要内容的前提下进行随机裁剪
颜色扰动（Color Distortion）：调整图像的亮度、对比度、饱和度和色调
高斯模糊（Gaussian Blur）：对图像应用轻微的高斯模糊

这些增强的组合确保了正样本对在视觉上相似但又不完全相同，迫使模型学习到更鲁棒的特征。

投影头（Projection Head）是 SimCLR 框架中的重要组件，它是一个小型神经网络，将编码器提取的特征映射到对比损失计算的空间。典型设计包括：

一个全连接层（输入维度 = 编码器输出维度，输出维度 = 投影维度）
ReLU 激活函数
另一个全连接层（输出维度 = 最终投影维度）

投影头的作用是将特征转换到更适合计算对比损失的空间，训练完成后可以丢弃，只保留编码器用于下游任务。

温度系数（Temperature Parameter）τ 在对比损失中起到重要作用。数学上，对比损失（NT-Xent loss）定义为：

L = -log[exp(sim(z_i, z_j)/τ) / Σ exp(sim(z_i, z_k)/τ)]

其中：
– sim() 是余弦相似度函数
– z_i 和 z_j 是正样本对
– z_k 是负样本
– τ 控制着正负样本的区分强度

τ 值较小时，模型会更关注最难区分的负样本；τ 值较大时，模型对所有负样本的关注更均匀。通常 τ 在 0.05 到 0.2 之间效果最佳。

以下是基于 CIFAR-10 数据集的 SimCLR 实现代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, datasets

# 数据增强
train_transform = transforms.Compose([transforms.RandomResizedCrop(32),
    transforms.RandomHorizontalFlip(),
    transforms.RandomApply([transforms.ColorJitter(0.8, 0.8, 0.8, 0.2)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.GaussianBlur(kernel_size=int(0.1*32)),
    transforms.ToTensor(),])

# 加载 CIFAR-10 数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=4)

# 编码器网络（使用 ResNet18 作为示例）class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=False)
        self.net.fc = nn.Identity()  # 移除最后的全连接层

    def forward(self, x):
        return self.net(x)

# 投影头
class ProjectionHead(nn.Module):
    def __init__(self, input_dim=512, hidden_dim=256, output_dim=128):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# SimCLR 模型
class SimCLR(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.projection = ProjectionHead()

    def forward(self, x):
        features = self.encoder(x)
        projections = self.projection(features)
        return F.normalize(projections, dim=1)

# 对比损失
def contrastive_loss(projections, temperature=0.1):
    device = projections.device
    batch_size = projections.shape[0] // 2

    # 计算相似度矩阵
    sim_matrix = torch.mm(projections, projections.t().contiguous())
    sim_matrix = sim_matrix / temperature

    # 创建标签：每个样本的正样本是它的增强版本
    labels = torch.arange(batch_size, device=device)
    labels = torch.cat([labels, labels])

    # 计算交叉熵损失
    loss = F.cross_entropy(sim_matrix, labels)
    return loss

# 训练循环
model = SimCLR().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)

for epoch in range(100):
    for batch, _ in train_loader:
        batch = batch.cuda()
        # 获取增强视图
        x1, x2 = batch.chunk(2)
        # 前向传播
        z1 = model(x1)
        z2 = model(x2)
        projections = torch.cat([z1, z2], dim=0)
        # 计算损失
        loss = contrastive_loss(projections)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch}, Loss: {loss.item()}')

对比学习通常需要较大的 batch size（256-4096）来提供足够的负样本。但当 GPU 显存有限时，可以：