最近有个项目要做图像分类,手头有两万张完全没标注的图片,人工标注成本实在顶不住。我想着能不能让模型自己从数据里学点有用的表征,再拿少量标注去做微调。参考了一圈之后,决定用PyTorch搭一套自监督预训练流程,跑了两轮实验就把精度追到了接近全监督的水平。这篇东西就来聊聊,怎么用PyTorch把自监督学习快速落地,从环境准备到核心原理,再到代码实现和显存优化,把能踩的坑先帮你扫一遍。
如果你是刚接触深度学习的同学,想搞清楚自监督学习到底怎么玩,或者已经在做有监督任务、想把手头无标注数据利用起来,这篇文章应该能帮你在最短时间内搭出一条能跑通的完整路线。我尽量不堆术语,遇到绕不开的概念会用最直白的方式解释,同时把每一步的“为什么这么做”讲清楚。
1. 自监督学习与PyTorch:为什么这个组合越来越常见
1.1 自监督学习解决什么问题
先聊一个很现实的问题:深度模型强归强,但它本质上是“数据驱动”的,你喂它多少标注,它就还你多少精度。但标注这件事,成本高得离谱。医院影像数据要请专业医生标,工业缺陷数据要请老师傅标,哪怕只是网上下载的图片打标签,也够你烦一天的。尤其是垂直场景,数据可能就几千张,标注完了还不够模型塞牙缝。
自监督学习的思路很直接:不依赖人工标签,让模型从数据自身寻找“监督信号”。比如把一张图片随机遮挡一部分,让模型去预测被遮住的内容;或者把图片做一些增强变换,让模型学会辨认同一个物体的不同形态。模型做完这些“自造任务”之后,会积累大量关于数据分布和特征的知识,这些知识再迁移到下游小样本任务上,效果好得惊人。
说白了,自监督相当于让模型先自学一遍打通任督二脉,再用极少量的高质量标注去“指点”一下,就达到甚至超过从头训练的效果。就我实测来看,在只有几百张标注图的情况下,先用自监督做预训练,精度比直接有监督训练能高出十几个百分点,这个优势在数据越少时越明显。
1.2 为什么用PyTorch来做这件事
市面上深度学习框架不少,但做自监督研究的人大多集中在PyTorch这边。原因说起来也不复杂。
首先,PyTorch的动态图机制对写自监督实验特别友好。自监督的很多方法需要在线修改数据流、动态构造正负样本对、在不同分支之间共享权重,这些操作在动态图下写起来像写普通Python一样自然。你调试的时候可以在loss计算之前随意插入print,随时看中间变量,这种亲民感是早期静态图框架给不了的。
其次,PyTorch生态里现成的模型库、预训练权重、第三方实现都很齐全。以自监督里最常用的backbone——ResNet和ViT为例,torchvision几行代码就能加载,huggingface上也有各种预训练版本。很多前沿方法的官方代码是PyTorch写的,你要复现就省去了一轮翻译的功夫。
再有一个很实际的原因:PyTorch对显存的控制更透明。自监督训练往往需要大批量才能出效果,显存优化是绕不过去的一关。PyTorch的混合精度、梯度累积、激活检查点这些机制用起来都很顺手,后面我会逐一讲怎么靠这些手段把一张卡吃满。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 环境与工程准备:开工前需要确认的三件事
2.1 Python版本与PyTorch安装包的选择
可能有人觉得环境配置不是技术活,随便装一下就行,但实际上自监督学习对版本兼容性异常敏感。我自己就吃过亏:一次在Windows上装好了PyTorch 1.13,写代码时用到torch.utils.data.Dataset的新接口,结果和旧版Python的行为不一致,排查了半天才发现是版本混搭的问题。
先说结论,目前跑自监督实验我推荐这套组合:Python 3.9或3.10 + PyTorch 2.x + CUDA 11.8或12.1。Python版本不用追求最新,3.12虽然新,但有些库的wheel还没跟上,容易遇到“装不上、跑不起”的尴尬。PyTorch 2.x带来了torch.compile加速,实测在部分模型上能提升20%-30%的训练速度,这对“超快”这个目标很有价值。
安装的时候优先用官方推荐的pip命令,PyTorch官网首页会根据你的系统和CUDA版本生成对应的安装指令。如果服务器上已经有CUDA,可以通过nvidia-smi查看驱动支持的最高CUDA版本,然后向下兼容安装PyTorch对应的CUDA版本即可。注意,PyTorch安装包里的CUDA是运行时依赖,你不需要单独装一套完整CUDA Toolkit,驱动匹配就行。
2.2 GPU驱动、CUDA与PyTorch的版本匹配
这块是新手最容易翻车的点。很多人跑起来显示torch.cuda.is_available()返回False,心态直接崩掉。实际上大多数情况不是代码问题,而是驱动、CUDA、PyTorch三者之间的版本组合出了问题。
我一般用这个顺序检查:先看驱动版本,nvidia-smi右上角显示的是驱动支持的最高CUDA版本;再看PyTorch自带的CUDA版本,用torch.version.cuda查看;最后确认PyTorch版本,用torch.__version__查看。理论上要求驱动的最高CUDA版本大于等于PyTorch编译时的CUDA版本即可。
举个例子,显卡驱动支持CUDA 12.2,那么装cu121或cu118的PyTorch都能正常跑。反过来,如果装了cu124的PyTorch而驱动只支持CUDA 12.0,就会报CUDA driver version is insufficient。判断方法很简单,安装前先搞清楚自己驱动支持什么版本,再去官网选对应安装包,省掉无数烦恼。
2.3 下载慢的处理方案:镜像源与离线包
PyTorch安装包动辄几百MB甚至2GB,从官方源下载慢到怀疑人生。最开始我按网上的教程用默认源装,下到一半断连,重来好几遍,心累得很。
这里分享一个我踩坑之后学乖的方案。下载PyTorch本体时用国内镜像源,比如清华源或阿里源,把conda或pip的下载地址切过去就行。但需要注意,PyTorch官网安装命令里有个--index-url参数,指的是PyTorch自己的源,你在后面加-i国内镜像源时,两者会冲突。我的做法是:先把--index-url里的链接下载成.whl文件,然后再用本地文件安装。
具体来说,在浏览器或下载工具里打开https://download.pytorch.org/whl/cu118/torch-2.1.0%2Bcu118-cp39-cp39-win_amd64.whl这样的链接,下载完成后在命令行执行pip install 本地路径。这样不仅快,还方便你保存一个备份,之后给同事或换机器时直接复用。实测下来,在下载高峰期用这个方式能节省大量时间。
3. 自监督学习的核心技术路径拆解
3.1 对比学习:让模型学会“找相同”
对比学习是目前自监督学习里声势最浩大的一支,代表方法包括SimCLR、MoCo、BYOL。它的核心思想可以概括成一句话:让模型学会把同一个样本的不同视角拉近,把不同样本的视角推远。
“视角”这个词听起来玄乎,其实就来自数据增强。比如同一张猫的照片,随机裁剪出一块、做颜色扰动、旋转一下,得到两个看起来不完全一样但依然是猫的新图片,这就是两个正样本视角。训练时,模型把这两张图分别编码成向量,然后优化目标希望这两个向量的相似度尽可能高,同时与批次里其他图片的向量相似度尽可能低。
这个思路有一个让我印象很深的理解角度:对比学习其实是在教模型建立一种“不变量”的感知能力。不管猫趴在沙发上还是躺在椅子上,不管光线是亮是暗,模型都要把核心的“猫”的特征提取出来,而忽略光线、背景等次要因素。这种能力迁移到下游分类、检测、分割任务中非常好用。
3.2 掩码重建:让模型学会“补全”
另一条技术路线是掩码重建,最出名的代表是MAE(Masked Autoencoder)。这个思路的灵感来自语言模型中的完形填空:你给模型一张被随机遮挡掉大部分区域的图片,让它根据剩余部分还原整张图。模型为了完成重建任务,被迫学到丰富的语义信息。
MAE最“激进”的地方在于,它会随机遮挡图片中75%的Patch,只保留25%给编码器看。这样做的目的是强迫模型不能只靠局部纹理和颜色猜测,而必须理解整体结构。比如看到一只猫的耳朵和尾巴,就要推断出中间大概率是身体,这种推断能力本质上就是高层次的语义理解。
在PyTorch里实现MAE比想象中简单,核心就是把图片切成Patch、随机mask掉一部分、用Transformer编码可见部分、再用解码器重建全部像素。不过MAE对显存要求比较高,通常要配合ViT-Base或ViT-Large使用,我在单张24GB显卡上跑过ViT-Base,batch size只能开到32左右,这时就需要后面的优化技巧来救场。
3.3 我应该选哪一种范式
很多朋友喜欢问“对比学习和掩码重建哪个更好”,这其实没有标准答案,完全取决于你手头的资源和数据。
对比学习对backbone的约束更少,ResNet、ViT都能用,训练也更加稳定,而且对batch size的容忍度较高。SimCLR原文用了巨大的batch size(8192),但我在小规模任务上实测,batch size设为256也能学到不错的特征。如果你打算做中小规模图像任务,对比学习是更省心的选择。
掩码重建在语义理解上潜力更大,尤其适合大规模预训练,但它对计算资源的要求也更高。ViT的引入基本是必须的,训练轮次更长,收敛速度相对较慢。如果算力有限但数据量又很多,我建议优先上MAE;如果只是想快速提升现有任务的效果,对比学习上手更快。
我的个人经验是:先用对比学习搭一套基线,跑通流程拿到下游任务结果,如果指标不够好再切换或融合掩码重建。这样风险最低,也不会一头扎进大模型的坑里出不来。
4. PyTorch快速实现一个自监督预训练流程(SimCLR精简版)
4.1 数据增强:两个views的构造
这里我直接以SimCLR为例,给出一套可以在CIFAR-10上运行的精简实现。整套代码逻辑完整,跑通之后替换成自己的数据集就能用。
SimCLR的核心是对每个样本生成两个不同的增强视图。在PyTorch里,我通过torchvision.transforms把增强算子叠成一个组合,然后在取数据时对一个样本应用两次。实际代码里,我会把增强流程定义成一个SimCLRTransform类,方便被DataLoader调用。
python复制import torch
import torchvision
import torchvision.transforms as T
class SimCLRTransform:
def __init__(self, size=32):
self.transform = T.Compose([
T.RandomResizedCrop(size=size, scale=(0.2, 1.0)),
T.RandomHorizontalFlip(),
T.RandomApply([T.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
T.RandomGrayscale(p=0.2),
T.ToTensor(),
T.Normalize(mean=[0.4914, 0.4822, 0.4465],
std=[0.2023, 0.1994, 0.2010])
])
def __call__(self, x):
return self.transform(x), self.transform(x)
增强策略的选择不是随便来的。SimCLR原文专门做过实验,发现随机裁剪和颜色扰动是对比学习效果最强的两个增强方式。裁剪让模型关注不同尺度的局部特征,颜色扰动则防止模型偷懒地依赖颜色分布来判断相似性。
4.2 模型结构:Encoder与Projection Head
SimCLR的模型结构分两部分:主干编码器(Encoder)和投影头(Projection Head)。编码器负责把图片变成特征向量,投影头则把特征映射到一个更紧凑的对比空间。训练完成后,投影头会被丢弃,只保留编码器用于下游任务。这个细节非常关键,好多人第一次看代码会疑惑“训练完了怎么结构不一样了”。
python复制import torch.nn as nn
class ProjectionHead(nn.Module):
def __init__(self, in_dim=512, hidden_dim=2048, out_dim=256):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True),
nn.Linear(hidden_dim, out_dim)
)
def forward(self, x):
return self.net(x)
class SimCLRModel(nn.Module):
def __init__(self, base_encoder=torchvision.models.resnet18, projection_dim=256):
super().__init__()
self.encoder = base_encoder(num_classes=4) # 只取4维特征
in_dim = 4
self.projection = ProjectionHead(in_dim, 2048, projection_dim)
def forward(self, x):
h = self.encoder(x)
z = self.projection(h)
return h, z
代码里有个小细节需要注意:ResNet18的fc层我改成了输出4维,这样整个特征向量维度很小,后面投影头计算开销也小。如果你用更大的backbone,比如ResNet50,就把fc改成512或2048维,投影头的in_dim也要同步改。这一点我看很多人会漏掉,导致维度对不上报错。
4.3 对比损失:InfoNCE的向量化实现
对比损失函数是自监督训练的心脏,SimCLR用的是InfoNCE损失。它的计算公式不算复杂,但实现时写法可以很优雅。我把整个batch的增强视图拼在一起,通过矩阵乘法计算所有样本两两之间的相似度,然后用对角线值作为正样本相似度。
python复制def info_nce_loss(z, batch_size, temperature=0.07):
z = torch.nn.functional.normalize(z, dim=1)
similarity_matrix = torch.matmul(z, z.T)
mask = torch.eye(batch_size, dtype=torch.bool).to(z.device)
positives = similarity_matrix[mask].view(batch_size, 2)
positives = torch.cat([positives[:, 0].unsqueeze(1), positives[:, 1].unsqueeze(1)], dim=1)
negatives_mask = ~mask
negatives = similarity_matrix[negatives_mask].view(batch_size, -1)
logits = torch.cat([positives, negatives], dim=1)
labels = torch.zeros(logits.shape[0], dtype=torch.long).to(z.device)
loss = torch.nn.functional.cross_entropy(logits / temperature, labels)
return loss
写这个函数的时候,我特别注意了相似度矩阵的构建方式。similarity_matrix是2N×2N的矩阵,positives取的是每个样本和它自己增强视图的相似度,位于矩阵的对角线附近。在我的实现里,正样本相似度被拼到每行的前两个位置,后面跟着和其他所有负样本的相似度,然后交给交叉熵损失,让模型学习把正样本放在第一位。
温度系数temperature=0.07是原文调过参的,它控制着相似度分布的“锐度”。温度越低,模型会越激进地拉近正样本、推开负样本,但太低又容易导致训练不稳定。新手直接抄0.07就行,不用纠结。
4.4 训练主循环与模型保存
训练主循环和普通有监督训练没太大区别,要注意的是需要手动控制数据增强生成两个视图的顺序。我在生成DataLoader的时候会用一个简单的包装类,保证每次拿到的是一个包含两个视图的元组。
python复制class TwoViewsWrapper(torch.utils.data.Dataset):
def __init__(self, dataset):
self.dataset = dataset
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx):
img, _ = self.dataset[idx]
x_i, x_j = SimCLRTransform()(img)
return x_i, x_j
python复制from torch.utils.data import DataLoader
# 以CIFAR-10为例
base_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=torchvision.transforms.ToTensor())
train_dataset = TwoViewsWrapper(base_dataset)
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=4, drop_last=True)
model = SimCLRModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
for epoch in range(200):
epoch_loss = 0.0
for x_i, x_j in train_loader:
x_i, x_j = x_i.cuda(), x_j.cuda()
z_i = model(x_i)[1]
z_j = model(x_j)[1]
z = torch.cat([z_i, z_j], dim=0)
loss = info_nce_loss(z, batch_size=x_i.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
scheduler.step()
if epoch % 20 == 0:
torch.save(model.state_dict(), f'simclr_epoch{epoch}.pth')
这里有一点非常重要:info_nce_loss里的batch_size参数传的是x_i.size(0),也就是当前batch里单个视图的样本数。因为z是[x_i, x_j]拼接出来的,总行数是2N,所以函数里mask用的是batch_size而不是z.size(0),这点搞错的话整个损失就废了。
训练结束后,我用torch.save保存了模型权重。但正式使用的时候,只需要保留model.encoder部分,投影头的参数可以直接丢弃。为了实用,我通常会单独存一份只含Encoder权重的state_dict,方便后续加载做下游微调。
5. 把训练速度“超快”落地的几个技巧
5.1 混合精度:一张卡当两张卡用
自监督训练动辄几百个epoch,时间成本很高。我第一次跑SimCLR在单张2080Ti上,200个epoch花了将近两天,当时差点劝退。后来加了混合精度训练(AMP),速度提升了接近一倍,显存占用也降低了不少。AMP的思路很简单:在训练过程中用FP16做前向和反向计算,用FP32保存主权重和优化器状态,既节省显存又利用GPU的Tensor Core加速。
PyTorch里实现AMP非常简单,只需要改几行代码。
python复制scaler = torch.cuda.amp.GradScaler()
for x_i, x_j in train_loader:
x_i, x_j = x_i.cuda(), x_j.cuda()
with torch.cuda.amp.autocast():
z_i = model(x_i)[1]
z_j = model(x_j)[1]
z = torch.cat([z_i, z_j], dim=0)
loss = info_nce_loss(z, batch_size=x_i.size(0))
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
autocast负责把模型计算中的一部分算子自动切换成FP16,GradScaler负责把Loss放大一定倍数,防止梯度下溢到0。这两个模块搭配起来几乎不用关心底层细节,我每次开新实验都会优先把AMP加上。
不过AMP也不是完全没坑,个别算子在FP16下会有精度问题,比如BatchNorm在FP16下运行偶尔会导致loss抖动。如果遇到这种情况,可以把模型传给autocast前用model = model.float()强制主权重保持FP32,或者查阅PyTorch文档中关于fp16不支持算子的说明。
5.2 DataLoader加速:num_workers与prefetch机制
很多人训练慢,问题不在GPU,而在CPU喂数据的速度跟不上。自监督学习的增强操作比普通训练复杂得多,每次取数据都要做随机裁剪、颜色抖动,计算密集,如果DataLoader的num_workers设置太小,GPU会一直空转等数据,利用率惨不忍睹。
我建议把num_workers设为你CPU核心数的两倍左右,同时打开persistent_workers=True,避免每个epoch重复创建worker进程的开销。还有一个容易被忽略的参数是prefetch_factor,设置成2或4,可以让DataLoader提前预取多个batch的数据,进一步掩盖数据加载延迟。
python复制train_loader = DataLoader(
train_dataset,
batch_size=256,
shuffle=True,
num_workers=8,
drop_last=True,
persistent_workers=True,
prefetch_factor=4
)
我实测在8核CPU的机器上,num_workers从4调到8之后,GPU利用率从不到60%升到了90%以上。当然,num_workers也不是越大越好,worker太多会导致系统资源竞争,反而拖慢整体速度。你可以根据训练时nvidia-smi显示的GPU利用率来调整:如果利用率持续低于80%,优先加worker;如果加完反而变慢了,说明已经过犹不及。
5.3 先小规模验证,再全量训练
这条经验看着像废话,但确实是让我少走最多弯路的一招。自监督训练本身没有标签可以实时评估效果,你不可能靠跑一个batch就看出来模型学得好不好。如果一上来就直接全量数据训练,跑了50个epoch才发现loss爆了或者参数设置错了,几个小时就白白浪费了。
我的做法是:先用1-2个batch的数据跑一遍完整的训练循环,确认前向、反向、优化器更新没有任何问题,且loss能正常下降,再开一个小规模的“冒烟测试”。具体来说,从全量数据里抽1%的样本,比如CIFAR-10就从5000张里抽50张图,batch_size设小一点,运行5-10个epoch,观察loss曲线和显存占用。
如果小规模测试通过,再切换到全量数据。这个流程听着多花了一点时间,实际上能帮你避免大量返工。我见过太多人把batch_size设爆显存、learning rate设太大导致loss发散、增强写错导致模型学到噪声,这些问题在小规模测试中都会提前暴露。
6. 踩坑与排查:自监督训练常见问题的处理记录
6.1 loss不降或者乱降
自监督训练最常见的困惑就是loss曲线表现异常。一种情况是loss从头到尾不怎么下降,这通常意味着模型根本没有学到有效特征。我遇到过几次,排来排去发现是学习率太高,模型在loss曲面上震荡。调整方式是把学习率从3e-4降到1e-4,同时加一个warmup阶段,让学习率慢慢爬升,曲线就老实了。
另一种相反的情况是loss前100个iter降得飞快,后面就原地踏步了。这也不正常。原因大概率是模型学到了“捷径”,比如只靠图片的平均颜色或者背景纹理就能区分正负样本,而没有真正理解语义。解决思路是增强数据增强的强度,比如提高裁剪比例下限、增大颜色扰动的幅度,强制模型去学更鲁棒的特征。
还有一点容易被忽略:InfoNCE里的温度系数。如果loss一直在2.5附近不下来,多半是温度太低导致梯度太小,可以试试把0.07调成0.1或0.2。反之,如果loss下降非常慢,就把温度调低一点。这个参数跟数据规模和batch size有关系,不是固定值。
6.2 GPU利用率上不去
GPU利用率低是训练加速的拦路虎,也是新手特别容易忽视的问题。我排查的时候一般按三个步骤走:先看数据加载是不是瓶颈,调大num_workers和prefetch_factor;再看训练循环里是不是有频繁的GPU-CPU数据拷贝,比如在循环里用了.item()或.numpy(),这些操作会强制同步,打断GPU流水线;最后看模型本身的计算量,如果是网络太浅或输入太小,GPU计算速度远快于数据加载,那利用率低反而是正常的。
在自监督训练里,开销最大的增强操作要在CPU端完成,很容易拖慢整体速度。如果num_workers已经加到很大利用率还是上不去,可以考虑把一部分增强操作放到GPU上实现,比如用kornia库在张量上做随机裁剪和颜色扰动。这个方法我还没在实际项目中落地过,但看到不少大公司的工程博客提过,属于进阶优化手段,普通场景做好DataLoader调优就够了。
6.3 显存溢出(OOM)的处理
自监督训练因为需要同时处理两个视图,显存占用大约是普通训练的两倍,OOM几乎是人人都要经历的。我第一次跑SimCLR时用ResNet50和batch size 128,直接爆了24GB显存,报错信息一串红,当时心态确实有点崩。
处理OOM有几个实用手段,按性价比排序:第一,减小batch size,这最简单,但注意batch太小会影响对比学习的负样本数量,效果会打折;第二,开梯度累积,用多个小batch的梯度叠加起来模拟大batch的效果,对对比学习特别有效;第三,开启激活检查点(activation checkpointing),以更多计算换更少显存,理论上可以在显存不变的情况下翻倍网络深度。
python复制from torch.utils.checkpoint import checkpoint
def forward_with_checkpoint(encoder, x):
return checkpoint(encoder, x, use_reentrant=False)
用torch.utils.checkpoint包装一下模型的前向计算,训练时中间激活值就不会全部保留,而是反向传播时重新算一遍。这个方法对ResNet和ViT都有效,代码改动只有一行。不过它会让训练时间增加20%左右,如果你显存还没到极限,建议优先用梯度累积。
python复制accumulation_steps = 4
optimizer.zero_grad()
for step, (x_i, x_j) in enumerate(train_loader):
loss = compute_loss(x_i, x_j)
loss = loss / accumulation_steps # 把loss缩放
loss.backward()
if (step + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
这段代码里有个细节,我在做梯度累积时会把每个小batch的loss除以accumulation_steps,这样累积4次之后的总梯度就相当于一个大batch的梯度,不会因为batch size翻倍而把学习率搞乱。这个技巧在多个对比学习项目中帮我省下了不少显存。
最后再分享一个我在几轮实验里摸索出来的小技巧:自监督训练完成之后,千万不要急着把Encoder直接拿去下游任务。先用下游任务的数据做一轮轻量的有监督微调,让模型的特征分布和任务本身的分布对齐,精度会再上一个台阶。微调的时候可以保持预训练的参数不动,只训练最后新加的分类头,这样训练速度非常快,在几百张标注数据的情况下,几十个epoch就能收敛。这个“预训练-微调”的套路,是我在项目里真正吃到甜头的环节,也让我相信自监督这套玩法在工业场景里能做到又快又稳。
