1. SimCLR基础框架实现
任务目标
实现SimCLR基础框架，包含随机裁剪/色彩扰动等数据增强、两层投影头（Projection Head）以及NT-Xent对比损失，并在CIFAR-10上训练一个简化版ResNet-18编码器。训练完成后，通过线性评估（冻结编码器，训练线性分类头）验证表征质量，以分类准确率作为核心指标。
环境依赖
需要 torch、torchvision、numpy。若使用GPU，需安装对应CUDA版本PyTorch。以下代码默认在单卡GPU或CPU上运行。
核心实现
以下代码完整实现了SimCLR训练流程：数据增强管道、编码器+投影头、NT-Xent损失函数、训练循环以及线性评估。将代码保存为 simclr_cifar10.py 后直接运行即可。
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=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, out_dim=128):
super().__init__()
# 使用CIFAR-10适配的简化ResNet-18（去掉首层大卷积）
self.conv = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
)
self.fc = nn.Linear(256, 512)
self.projection = nn.Sequential(
nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, out_dim)
)
def forward(self, x):
h = self.conv(x).flatten(1)
h = F.relu(self.fc(h))
return self.projection(h), h # 返回投影特征和表征
# ---------- 3. NT-Xent损失 ----------
def nt_xent_loss(z1, z2, temperature=0.5):
batch_size = z1.size(0)
z = torch.cat([z1, z2], dim=0) # 2N x D
z = F.normalize(z, dim=1)
sim = torch.mm(z, z.t()) / temperature # 2N x 2N
# 屏蔽对角线（自身与自身）
mask = torch.eye(2 * batch_size, device=z.device).bool()
sim.masked_fill_(mask, -1e9)
# 正样本对：i 与 i+batch_size（互为增强视图）
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 = 128
epochs = 50 # 小型模型训练50轮，约10分钟（GPU）
lr = 3e-4
# 数据加载（CIFAR-10，仅用训练集做无监督预训练）
base_transform = T.Compose([T.ToTensor(), T.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
train_set = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=base_transform)
# 为节省时间，使用前10000张子集（可自行调整）
train_subset = Subset(train_set, range(10000))
augment = SimCLRAugment()
# 包装数据集以返回两个增强视图
class PairDataset(torch.utils.data.Dataset):
def __init__(self, subset):
self.subset = subset
def __len__(self):
return len(self.subset)
def __getitem__(self, idx):
img, _ = self.subset[idx]
x1, x2 = augment(img)
return x1, x2
pair_set = PairDataset(train_subset)
loader = DataLoader(pair_set, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=True)
model = Encoder().to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)
# ---------- 5. 预训练循环 ----------
model.train()
for epoch in range(epochs):
total_loss = 0.0
for x1, x2 in loader:
x1, x2 = x1.to(device), x2.to(device)
z1, _ = model(x1)
z2, _ = model(x2)
loss = nt_xent_loss(z1, z2)
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(loader):.4f}")
# ---------- 6. 线性评估（冻结编码器） ----------
print("开始线性评估...")
model.eval()
# 冻结卷积层和fc层（不更新）
for param in model.parameters():
param.requires_grad = False
# 提取特征的数据集
class FeatureSet(torch.utils.data.Dataset):
def __init__(self, subset):
self.subset = subset
def __len__(self):
return len(self.subset)
def __getitem__(self, idx):
img, label = self.subset[idx]
with torch.no_grad():
_, h = model(img.unsqueeze(0).to(device))
return h.cpu().squeeze(0), label
# 使用完整训练集（50000张）评估，但为速度可只用10000张
eval_train = Subset(train_set, range(10000))
eval_test = torchvision.datasets.CIFAR10(root="./data", train=False, download=True, transform=base_transform)
eval_test = Subset(eval_test, range(2000))
train_feat = FeatureSet(eval_train)
test_feat = FeatureSet(eval_test)
train_loader = DataLoader(train_feat, batch_size=256, shuffle=True)
test_loader = DataLoader(test_feat, batch_size=256, shuffle=False)
# 线性分类头
linear = nn.Linear(512, 10).to(device)
opt_lin = optim.Adam(linear.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(20):
linear.train()
for h, y in train_loader:
h, y = h.to(device), y.to(device)
out = linear(h)
loss = loss_fn(out, y)
opt_lin.zero_grad()
loss.backward()
opt_lin.step()
# 测试
if (epoch+1) % 5 == 0:
linear.eval()
correct = 0
total = 0
with torch.no_grad():
for h, y in test_loader:
h, y = h.to(device), y.to(device)
out = linear(h)
pred = out.argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
acc = 100.0 * correct / total
print(f"Linear Epoch [{epoch+1}/20] Test Acc: {acc:.2f}%")
运行与验证
在终端执行：
python simclr_cifar10.py程序会自动下载CIFAR-10数据集到 ./data 目录。训练过程中每10轮打印一次NT-Xent损失值，损失应呈下降趋势（例如从约4.0降至1.5以下）。线性评估阶段每5轮打印一次测试准确率。
通过判定标准
- 预训练损失持续下降，最终低于2.0（使用10000张子集、50轮训练时）。
- 线性评估准确率在20轮内达到40%以上（使用10000张训练子集、2000张测试子集）。若使用完整50000张训练集，准确率可提升至60%以上。
- 若准确率远低于上述值，可检查增强强度是否过大（如ColorJitter参数）、学习率是否合适，或增加训练轮数。
常见问题
- 显存不足：降低
batch_size至64或32。 - 损失不下降：确认数据增强已生效（可打印增强后的图像检查），并确认NT-Xent实现中正样本对索引正确。
- 线性评估准确率低：增加预训练轮数至100轮，或使用完整训练集；也可尝试增大投影头输出维度（如256）。