PyTorch 深度学习实战（28）：对比学习（Contrastive Learning）与自监督表示学习

在上一篇文章中，我们探讨了扩散模型（Diffusion Models）在图像生成中的应用。本文将重点介绍 对比学习（Contrastive Learning），这是一种通过构建正负样本对来学习数据表征的自监督学习方法。我们将使用 PyTorch 实现一个简单的对比学习模型，并在 CIFAR-10 数据集上进行验证。

---

一、对比学习基础

对比学习的核心思想是通过最大化相似样本对的相似性，同时最小化不相似样本对的相似性。这种方法无需人工标注数据，即可学习到具有判别性的特征表示。

1. 对比学习的核心组件

• 数据增强：

  • 通过随机裁剪、颜色变换等操作生成同一图像的不同视图，构建正样本对。

• 编码器网络：

  • 将输入数据映射到低维特征空间（如 ResNet）。

• 投影头：

  • 将特征映射到对比学习空间（通常使用 MLP）。

• 对比损失函数：

  • 常用的 InfoNCE 损失函数，通过温度参数控制样本对的区分度。

2. 对比学习的数学原理

InfoNCE 损失函数定义为：

3. 对比学习的优势

• 无需标注数据：

  • 通过自监督方式学习通用特征表示。

• 特征可迁移性强：

  • 预训练的特征可用于下游分类、检测等任务。

• 鲁棒性高：

  • 对数据增强和噪声具有较好的适应性。

---

二、CIFAR-10 实战

我们使用 PyTorch 实现对比学习模型，并在 CIFAR-10 数据集上预训练特征编码器，最后通过线性评估验证特征质量。

1. 实现步骤

1. 定义数据增强策略

2. 构建编码器（ResNet-18）和投影头（MLP）

3. 实现 InfoNCE 损失函数

4. 预训练特征编码器

5. 冻结编码器，训练线性分类器评估特征

2. 代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm
import numpy as np
​
# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
​
# 修正的数据增强策略
class ContrastiveTransformations:
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.RandomResizedCrop(32, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
        ])
    
    def __call__(self, x):
        return self.transform(x), self.transform(x)
​
# CIFAR-10 数据集
train_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, 
    transform=ContrastiveTransformations())  # 使用自定义转换
​
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=512, shuffle=True, num_workers=4, pin_memory=True)
​
# 修正的编码器网络（ResNet-18）
class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torchvision.models.resnet18(pretrained=False)
        # 调整第一层卷积以适应CIFAR-10的32x32输入
        self.net.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.net.maxpool = nn.Identity()  # 移除初始的maxpool
        self.out_dim = self.net.fc.in_features
        self.net.fc = nn.Identity()  # 移除最后的全连接层
​
    def forward(self, x):
        # 确保输入是4D张量
        if x.dim() == 3:
            x = x.unsqueeze(0)
        return self.net(x)
​
# 投影头（MLP）
class ProjectionHead(nn.Module):
    def __init__(self, input_dim=512, hidden_dim=256, output_dim=128):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )
​
    def forward(self, x):
        # 处理可能的维度问题
        if x.dim() > 2:
            x = x.view(x.size(0), -1)
        return self.mlp(x)
​
# 对比学习模型
class ContrastiveModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder().to(device)
        self.projection = ProjectionHead(self.encoder.out_dim).to(device)
    
    def forward(self, x):
        # 确保输入是4D
        if x.dim() == 3:
            x = x.unsqueeze(0)
        features = self.encoder(x)
        projections = self.projection(features)
        return F.normalize(projections, dim=1)
​
# InfoNCE 损失函数
def info_nce_loss(query, key, temperature=0.1):
    batch_size = query.shape[0]
    # 计算相似度矩阵
    logits = torch.mm(query, key.T) / temperature
    # 对角线是正样本对
    labels = torch.arange(batch_size).to(device)
    return F.cross_entropy(logits, labels)
​
# 初始化模型和优化器
model = ContrastiveModel().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
​
# 训练函数
def train_contrastive():
    model.train()
    total_loss = 0
    for (x1, x2), _ in tqdm(train_loader, desc='Training'):
        x1, x2 = x1.to(device), x2.to(device)
        
        # 前向传播
        z1 = model(x1)
        z2 = model(x2)
        
        # 计算对称损失
        loss = (info_nce_loss(z1, z2) + info_nce_loss(z2, z1)) / 2
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    return total_loss / len(train_loader)
​
# 训练循环
for epoch in range(1, 201):
    loss = train_contrastive()
    scheduler.step()
    print(f'Epoch {epoch:3d} | Loss: {loss:.4f}')
​
# 保存编码器
torch.save(model.encoder.state_dict(), 'contrastive_encoder.pth')
​
# 线性评估准备
# 标准化验证数据集
val_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
​
class ContrastiveCIFAR10(torchvision.datasets.CIFAR10):
    def __getitem__(self, index):
        img, label = super().__getitem__(index)
        return img, label
​
# 训练集使用验证集的标准化变换
train_linear_dataset = ContrastiveCIFAR10(
    root='./data', train=True, download=True, transform=val_transform)
train_linear_loader = torch.utils.data.DataLoader(
    train_linear_dataset, batch_size=512, shuffle=True, num_workers=4)
​
# 验证集
val_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=val_transform)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=512, shuffle=False)
​
# 创建线性模型
linear_model = nn.Sequential(
    Encoder(),
    nn.Linear(512, 10)
).to(device)
linear_model[0].load_state_dict(torch.load('contrastive_encoder.pth', weights_only=True))
​
# 冻结编码器
for param in linear_model[0].parameters():
    param.requires_grad = False
​
# 训练参数
optimizer = optim.Adam(linear_model[1].parameters(), lr=3e-4)
criterion = nn.CrossEntropyLoss()
​
# 评估函数
def evaluate():
    linear_model.eval()
    correct = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = linear_model(images)
            correct += (outputs.argmax(1) == labels).sum().item()
    return 100 * correct / len(val_dataset)
​
# 训练循环
for epoch in range(1, 101):
    linear_model.train()
    for images, labels in train_linear_loader:
        images, labels = images.to(device), labels.to(device)
        
        outputs = linear_model(images)
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    acc = evaluate()
    print(f'Epoch {epoch:3d} | Val Acc: {acc:.2f}%')

---

三、代码解析

1. 数据增强：

   • 使用随机裁剪、颜色抖动等生成正样本对。

2. 模型架构：

   • 编码器使用 ResNet-18，投影头为两层 MLP。

3. 损失计算：

   • 对称计算两个增强视图的 InfoNCE 损失。

4. 线性评估：

   • 冻结预训练编码器，仅训练线性分类层评估特征质量。

---

四、运行结果

运行代码将得到以下典型输出：

Epoch   1 | Loss: 4.2162
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:15<00:00,  1.30it/s]
Epoch   2 | Loss: 2.6235
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.27it/s]
Epoch   3 | Loss: 2.0762
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:18<00:00,  1.25it/s]
Epoch   4 | Loss: 1.7587
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.27it/s]
Epoch   5 | Loss: 1.5706
...
Epoch 194 | Loss: 0.3861
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.27it/s]
Epoch 195 | Loss: 0.3858
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.26it/s]
Epoch 196 | Loss: 0.3879
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.27it/s]
Epoch 197 | Loss: 0.3895
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.26it/s]
Epoch 198 | Loss: 0.3833
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.26it/s]
Epoch 199 | Loss: 0.3923
Training: 100%|████████████████████████████████████████████████████| 98/98 [01:17<00:00,  1.27it/s]
Epoch 200 | Loss: 0.3881
​
Epoch   1 | Val Acc: 68.01%
Epoch   2 | Val Acc: 71.27%
Epoch   3 | Val Acc: 72.37%
Epoch   4 | Val Acc: 73.22%
Epoch   5 | Val Acc: 73.96%
...
Epoch  94 | Val Acc: 83.06%
Epoch  95 | Val Acc: 82.99%
Epoch  96 | Val Acc: 83.01%
Epoch  97 | Val Acc: 82.96%
Epoch  98 | Val Acc: 83.12%
Epoch  99 | Val Acc: 83.09%
Epoch 100 | Val Acc: 83.00%

---

五、总结

本文实现了基于对比学习的自监督表征学习框架，通过 InfoNCE 损失在 CIFAR-10 数据集上学习到具有判别性的特征表示。实验表明，仅使用线性分类器即可达到 82% 以上的验证准确率，证明了对比学习的有效性。

在下一篇文章中，我们将深入计算机视觉的核心任务之一------《目标检测与 YOLO 实战》，探讨如何利用深度学习实现高效的目标检测系统。