2. 视觉特征提取与对比学习预训练

任务说明

本案例将使用对比学习（Contrastive Learning）在CIFAR-10数据集上预训练一个轻量级视觉编码器（ResNet-18），使其提取的图像特征具有区分度——即同类图像特征距离近，异类图像特征距离远。预训练完成后，通过线性探测（Linear Probe）评估特征质量。

环境依赖

需要 torch、torchvision、numpy。若使用CUDA，需对应版本的PyTorch；CPU亦可运行（速度较慢）。

核心实现

以下代码完整实现了对比学习预训练流程：数据增强、编码器、投影头、InfoNCE损失函数、训练循环以及线性探测评估。代码可直接运行，默认使用CPU，若检测到CUDA则自动使用GPU。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as T
from torch.utils.data import DataLoader, Subset
import numpy as np

# ---------- 1. 数据增强（SimCLR风格） ----------
class SimCLRAugment:
    def __init__(self, size=32):
        self.transform = T.Compose([
            T.RandomResizedCrop(size, scale=(0.2, 1.0)),
            T.RandomHorizontalFlip(),
            T.ColorJitter(0.4, 0.4, 0.4, 0.1),
            T.RandomGrayscale(p=0.2),
            T.ToTensor(),
            T.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
    def __call__(self, x):
        return self.transform(x), self.transform(x)  # 返回两个增强视图

# ---------- 2. 编码器 + 投影头 ----------
class Encoder(nn.Module):
    def __init__(self, feat_dim=128):
        super().__init__()
        # 去掉ResNet18最后的全连接层，输出512维特征
        resnet = torchvision.models.resnet18(weights=None, num_classes=0)
        self.backbone = resnet
        self.proj = nn.Sequential(
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, feat_dim)
        )
    def forward(self, x):
        feat = self.backbone(x)  # [B, 512]
        return self.proj(feat)   # [B, feat_dim]

# ---------- 3. InfoNCE损失 ----------
def info_nce_loss(z, batch_size, temperature=0.5):
    # z: [2*B, D]，前B个是view1，后B个是view2
    z = F.normalize(z, dim=1)
    sim = torch.matmul(z, z.T) / temperature  # [2B, 2B]
    mask = torch.eye(2*batch_size, device=z.device).bool()
    sim.masked_fill_(mask, -1e9)
    # 正样本对：i 与 i+B（及对称）
    labels = torch.cat([
        torch.arange(batch_size, 2*batch_size),
        torch.arange(0, batch_size)
    ], dim=0).to(z.device)
    loss = F.cross_entropy(sim, labels)
    return loss

# ---------- 4. 训练配置 ----------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 256
epochs = 50
feat_dim = 128

# 数据加载（CIFAR-10，仅用训练集做预训练）
base_transform = T.Compose([T.ToTensor(), T.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
train_data = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=base_transform)
# 为对比学习使用双视图增强
train_data.transform = SimCLRAugment()
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=True)

model = Encoder(feat_dim).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# ---------- 5. 预训练循环 ----------
print("开始对比学习预训练...")
for epoch in range(epochs):
    total_loss = 0.0
    for (x1, x2), _ in train_loader:
        x1, x2 = x1.to(device), x2.to(device)
        z1 = model(x1)
        z2 = model(x2)
        z = torch.cat([z1, z2], dim=0)  # [2B, D]
        loss = info_nce_loss(z, batch_size)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    if (epoch+1) % 10 == 0:
        print(f"Epoch [{epoch+1}/{epochs}] Loss: {total_loss/len(train_loader):.4f}")

# ---------- 6. 线性探测评估 ----------
print("\n开始线性探测评估...")
# 冻结编码器，提取特征
model.eval()
feat_train, label_train = [], []
with torch.no_grad():
    for (x, _), labels in DataLoader(train_data, batch_size=256, shuffle=False):
        # 注意：train_data的transform是SimCLRAugment，这里需要重新用原始transform
        pass  # 下面重新加载数据

# 重新用原始transform加载数据用于评估
eval_transform = T.Compose([T.ToTensor(), T.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
eval_data = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=eval_transform)
eval_loader = DataLoader(eval_data, batch_size=256, shuffle=False)

with torch.no_grad():
    for x, labels in eval_loader:
        x = x.to(device)
        feat = model.backbone(x)  # 使用骨干网络特征（512维）
        feat_train.append(feat.cpu())
        label_train.append(labels)
feat_train = torch.cat(feat_train).numpy()
label_train = torch.cat(label_train).numpy()

# 训练线性分类器（逻辑回归）
from sklearn.linear_model import LogisticRegression
clf = LogisticRegression(max_iter=1000, C=1.0)
clf.fit(feat_train, label_train)

# 在测试集上评估
test_data = torchvision.datasets.CIFAR10(root="./data", train=False, download=True, transform=eval_transform)
test_loader = DataLoader(test_data, batch_size=256, shuffle=False)
feat_test, label_test = [], []
with torch.no_grad():
    for x, labels in test_loader:
        x = x.to(device)
        feat = model.backbone(x)
        feat_test.append(feat.cpu())
        label_test.append(labels)
feat_test = torch.cat(feat_test).numpy()
label_test = torch.cat(label_test).numpy()

acc = clf.score(feat_test, label_test)
print(f"线性探测准确率: {acc*100:.2f}%")

运行方式

安装依赖：pip install torch torchvision scikit-learn
将上述代码保存为 contrastive_pretrain.py
运行：python contrastive_pretrain.py

验证方法与通过标准

程序会输出每个epoch的损失值以及最终线性探测准确率。判定预训练有效的标准：

训练过程中InfoNCE损失持续下降（例如从初始约3.0降至0.5以下）；
线性探测准确率明显高于随机猜测（CIFAR-10随机为10%），通常应达到50%以上（训练50轮、无额外正则时合理预期为55%～65%）。

若准确率远低于此范围，可尝试增加训练轮数、调整学习率或增大batch size。

常见问题

显存不足：降低batch_size（如128或64），同时相应调整学习率。
训练速度慢：使用GPU可显著加速；若只有CPU，建议将epochs减少至20并观察损失下降趋势。
下载数据集失败：确保网络可访问CIFAR-10下载源，或手动下载后放入./data目录。

📚 课程目录