PyTorch自监督学习实战：从对比学习到掩码重建

最近有个项目要做图像分类，手头有两万张完全没标注的图片，人工标注成本实在顶不住。我想着能不能让模型自己从数据里学点有用的表征，再拿少量标注去做微调。参考了一圈之后，决定用PyTorch搭一套自监督预训练流程，跑了两轮实验就把精度追到了接近全监督的水平。这篇东西就来聊聊，怎么用PyTorch把自监督学习快速落地，从环境准备到核心原理，再到代码实现和显存优化，把能踩的坑先帮你扫一遍。

如果你是刚接触深度学习的同学，想搞清楚自监督学习到底怎么玩，或者已经在做有监督任务、想把手头无标注数据利用起来，这篇文章应该能帮你在最短时间内搭出一条能跑通的完整路线。我尽量不堆术语，遇到绕不开的概念会用最直白的方式解释，同时把每一步的“为什么这么做”讲清楚。

1. 自监督学习与PyTorch：为什么这个组合越来越常见

1.1 自监督学习解决什么问题

先聊一个很现实的问题：深度模型强归强，但它本质上是“数据驱动”的，你喂它多少标注，它就还你多少精度。但标注这件事，成本高得离谱。医院影像数据要请专业医生标，工业缺陷数据要请老师傅标，哪怕只是网上下载的图片打标签，也够你烦一天的。尤其是垂直场景，数据可能就几千张，标注完了还不够模型塞牙缝。

自监督学习的思路很直接：不依赖人工标签，让模型从数据自身寻找“监督信号”。比如把一张图片随机遮挡一部分，让模型去预测被遮住的内容；或者把图片做一些增强变换，让模型学会辨认同一个物体的不同形态。模型做完这些“自造任务”之后，会积累大量关于数据分布和特征的知识，这些知识再迁移到下游小样本任务上，效果好得惊人。

说白了，自监督相当于让模型先自学一遍打通任督二脉，再用极少量的高质量标注去“指点”一下，就达到甚至超过从头训练的效果。就我实测来看，在只有几百张标注图的情况下，先用自监督做预训练，精度比直接有监督训练能高出十几个百分点，这个优势在数据越少时越明显。

1.2 为什么用PyTorch来做这件事

市面上深度学习框架不少，但做自监督研究的人大多集中在PyTorch这边。原因说起来也不复杂。

首先，PyTorch的动态图机制对写自监督实验特别友好。自监督的很多方法需要在线修改数据流、动态构造正负样本对、在不同分支之间共享权重，这些操作在动态图下写起来像写普通Python一样自然。你调试的时候可以在loss计算之前随意插入print，随时看中间变量，这种亲民感是早期静态图框架给不了的。

其次，PyTorch生态里现成的模型库、预训练权重、第三方实现都很齐全。以自监督里最常用的backbone——ResNet和ViT为例，torchvision几行代码就能加载，huggingface上也有各种预训练版本。很多前沿方法的官方代码是PyTorch写的，你要复现就省去了一轮翻译的功夫。

再有一个很实际的原因：PyTorch对显存的控制更透明。自监督训练往往需要大批量才能出效果，显存优化是绕不过去的一关。PyTorch的混合精度、梯度累积、激活检查点这些机制用起来都很顺手，后面我会逐一讲怎么靠这些手段把一张卡吃满。

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文件，然后再用本地文件安装。

具体来说，在浏览器或下载工具里打开`http