1. 为什么所有CV入门教程都以torchvision为起点
先说一个我经常遇到的现象:很多朋友刚接触深度学习,看了一圈课程后还是不知道怎么把自己的图片数据跑起来。PyTorch其实是那种让你“灵活得不知道从哪下手”的框架,而torchvision恰恰补上了这个空档。它内部封装了大量计算机视觉开发中高频使用的工具——从数据集下载、数据预处理、经典模型到工具函数,几乎覆盖了一个视觉项目从数据到模型再到训练评估的全链路。我自己带过不少实习生,给他们的第一个任务从来不是写什么自定义网络,而是先让他们把torchvision的datasets、transforms和models三个模块彻底用熟。原因很简单:一个项目里90%的代码都在处理数据接入和变换,训练代码反而是最模式化的那部分。
torchvision到底解决了什么问题?说白了,它让“把一张普通图片变成模型能吃的张量”这件事变得标准化。你可以直接使用它内置的CIFAR10、ImageNet、COCO这些经典数据集来验证想法,也可以把日常积累的图片文件夹通过ImageFolder一行代码加载成Dataset。更关键的是,transforms模块把随机裁剪、翻转、归一化、颜色抖动这些数据增强操作封装成了可以叠加的“乐高积木”,你只需要用Compose把它们组合起来,就能快速构造出一套接近真实场景的输入分布。
有人可能会问:既然PyTorch本身已经这么强大了,为什么还要学torchvision?我的回答是:PyTorch提供一个灵活的神经网络框架,但计算机视觉领域有很多约定俗成的数据规范和预处理算法,torchvision把这些重复劳动沉淀成了稳定且经过社区验证的API。比如它的标准化参数,ImageNet的mean和std是长期统计出来的,你用这些参数去归一化预训练模型能获得更好的迁移效果——这背后有大量实验支撑,轮子不必自己造。
从我个人的学习路径看,熟悉torchvision的过程就是熟悉计算机视觉项目通用流程的过程。先搞懂数据从哪来、怎么进模型,再谈网络结构设计和训练调参,是一个性价比极高的路径。这篇内容不会去复述官方文档的每一个函数签名,而是把我在真实项目中反复用到的核心模块和进阶技巧拆开讲清楚,并且会把一些只有踩过坑才能总结出来的经验一并放进去。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 安装与版本匹配:那些404和下载慢的真相
2.1 不要只盯着官方pip源
提到安装torchvision,很多人的第一反应就是执行pip install torchvision。这条命令本身没错,但如果你在国内网络环境下执行,大概率会卡在下载阶段——官方源的PyTorch轮子普遍有200MB以上,再加上依赖的torch和torchaudio,动辄几个GB的流量会让人等得抓狂。我见过不少人因为下载速度慢到怀疑人生,甚至中途开手机热点硬扛,结果快下完时网络一抖就前功尽弃。热搜里就有人问“安装pytorch时,手机开了热点,下载速度还是很慢,怎么办”,这说明这问题太普遍了。
最好的办法是配置国内镜像源。清华的PyPI镜像、阿里云的PyPI镜像,这些不光是稳定,而且支持很多历史版本的torchvision,完全能满足日常开发需求。使用方法很简单:
bash复制pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
用镜像源之后,速度基本能从几十KB/s跑到几MB/s甚至更高。不过有一点要提醒:镜像源上的版本同步可能会有延迟,如果官方刚发布新版,镜像还没同步过来,这时候可以指定一个稍早的版本,或者干脆用官方源但配合下载工具。
2.2 CUDA版本与torch/torchvision的对应关系
安装GPU版PyTorch比CPU版麻烦的地方在于CUDA版本匹配。很多人一上来就装最新版,结果发现torchvision的版本和torch版本对不上,导致运行时import torchvision直接抛错。实际上,torchvision和torch的版本号是有严格对应关系的,不能随便组合。比如torch==2.5.1对应torchvision==0.20.1、torchaudio==2.5.1,这是官方在发布时约定好的。
以CUDA 11.8为例,如果你用官方源安装,命令是这样的:
bash复制pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu118
如果你的显卡驱动较新,支持更高版本CUDA,也可以选择对应CUDA 12.x的轮子。这里的关键是先确认自己机器的CUDA环境。可以在终端执行:
bash复制nvidia-smi
这个命令会输出GPU驱动版本和驱动支持的CUDA版本。注意这里的CUDA版本是驱动支持的,而不是你本地安装的CUDA Toolkit版本。PyTorch的CUDA轮子自带运行时,所以只要驱动支持某个CUDA版本,你就可以直接装对应版本号的PyTorch轮子,不需要额外安装CUDA Toolkit。
我曾经在AMD平台上尝试安装CUDA版PyTorch,这里有个容易混淆的点:AMD GPU不能直接用CUDA,你需要用ROCm版本或者通过DirectML等桥接方案。不过大多数情况,AMD平台可以用CPU版本的PyTorch跑小模型做学习实验,性能也可以接受。如果是稍大一点的模型,建议还是使用NVIDIA显卡。
2.3 手动解决mnist等数据集404问题
torchvision的datasets模块会自动下载MNIST等数据集,但在某些时刻,由于官方数据托管路径变更或网络问题,你会遇到下载时返回404的错误。比如torchvision.datasets.MNIST在下载训练集时可能会请求http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz,这个地址的可用性并不总是稳定。
遇到404时,比较好的方案是手动把数据集下载到本地。以MNIST为例,你可以去第三方镜像或直接搜索可用的下载链接,将压缩包放入./data/MNIST/raw目录下(具体目录结构取决于你的Dataset构造函数参数)。然后设置download=False重新加载,torchvision就会直接读取本地文件,不再触发网络请求。
还有一个思路是从内网或国内数据源下载。例如,一些高校或者云厂商会同步公开数据集,搜索引擎能帮你找到镜像。我自己更推荐把常用数据集提前下载并归档到本地磁盘,然后写一个简单的脚本检测文件是否存在,避免每次新建环境都要重新下载。这个习惯能节省大量时间,而且也让实验可复现性更强。
3. datasets与transforms:数据管线的核心拼图
3.1 Dataset到底在做什么?别只当它是“装载器”
很多初学者把torchvision.datasets理解成“一个用来取数据的列表”,这个理解太浅了。其实Dataset的核心是把“数据文件的路径/字节内容”和“模型输入张量”建立映射。它不只是返回一张图片,而是返回(样本, 标签)这样一个元组,而且允许你在__getitem__里嵌入任何所需的预处理逻辑。
torchvision内置了种类繁多的Dataset实现。常见的包括:
ImageFolder:适用于根目录下按类别组织图片文件夹的场景,非常方便;CIFAR10/100、MNIST、FashionMNIST:经典小规模学术数据集;COCO、VOCDetection:目标检测、分割任务的数据集;CelebA:人脸属性数据集。
以ImageFolder为例,假设你的数据目录结构是:
code复制data/train/
cat/
cat_001.jpg
cat_002.jpg
dog/
dog_001.jpg
那么使用方式:
python复制from torchvision import datasets
dataset = datasets.ImageFolder(root='data/train')
它会自动扫描子文件夹,并为每个类别分配一个数值标签。注意,ImageFolder并不限制图片格式,jpg、png都可以,但要求图片本身可以被PIL打开。我第一次用这个类的时候没注意某几张图是损坏的,结果训练到一半直接报错。后来学乖了,先跑一个扫描脚本把损坏图片剔除。
3.2 transforms是怎么一步步把图片变成张量的
torchvision.transforms是数据预处理的核心。它的设计理念是“函数式组合”而不是面向对象的调用。你通过transforms.Compose把多个操作串联起来,模型在每次取样本时会依次执行这些操作。
一个最典型的图像分类预处理管道是:
python复制from torchvision import transforms
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
这段代码做了四件事:统一图片尺寸、裁出中心区域、转为张量、按通道归一化。每一步单独看都不难,但合在一起就是一套标准的ImageNet风格输入预处理。如果你打算微调torchvision预训练模型,这套参数几乎就是一个“金标准”。
关于ToTensor,有个细节容易被忽略:它会将PIL图片或NumPy数组从[H, W, C]的uint8格式(范围0~255)转换为[C, H, W]的float32格式(范围0.0~1.0)。归一化基于这个0到1的范围进行,所以mean和std的值也都在0~1之间。如果你需要做特殊归一化,比如把数据范围映射到[-1, 1],那么可以直接用transforms.Normalize((0.5,), (0.5,)),这相当于(x - 0.5) / 0.5。
3.3 自定义Dataset的标准姿势
内置Dataset再好用,也不可能覆盖千奇百怪的实际场景。大部分真实项目都需要写一个自定义Dataset类。常写的场景包括:图片路径从CSV读取、样本带有多标签、需要同时返回mask等。
一个典型的自定义Dataset骨架如下:
python复制import os
from PIL import Image
from torch.utils.data import Dataset
class MyDataset(Dataset):
def __init__(self, img_dir, label_file, transform=None):
self.img_paths = []
self.labels = []
with open(label_file, 'r') as f:
for line in f:
filename, label = line.strip().split(',')
self.img_paths.append(os.path.join(img_dir, filename))
self.labels.append(int(label))
self.transform = transform
def __len__(self):
return len(self.img_paths)
def __getitem__(self, idx):
img = Image.open(self.img_paths[idx]).convert('RGB')
if self.transform:
img = self.transform(img)
return img, self.labels[idx]
这里有两个经验:一是Image.open之后最好调用convert('RGB'),防止某些灰度图只有一个通道导致后续张量维度对不上。二是__getitem__里的transform操作最好和加载分离,这样同一个Dataset可以配合不同的transform复用,比如训练时用增强,验证时只用基础变换。
3.4 数据加载器DataLoader的线程与随机性设置
Dataset定义了单个样本的取法,DataLoader则负责批量化、打乱、多进程预取。这里有几个很多人忽略的细节。
首先,num_workers并非越大越好。在Windows上,这个参数要小心使用,因为Windows的spawn启动方式比Linux的fork开销大,设置太高容易导致内存爆炸或启动缓慢。一个稳妥的做法是:在Linux服务器上可以设置成CPU核心数的一半,在Windows上可以先从0或1开始,再逐步上调。
其次是shuffle参数。训练集需要打乱,验证集不应该打乱(除非你想模拟随机效果)。如果你用分布式训练,还需要设置sampler而不是直接依赖DataLoader的shuffle。
再就是pin_memory。当使用GPU训练时,设置pin_memory=True可以让数据从内存传送到GPU显存时更快。因为pinned memory使用页锁定内存,会减少一次DMA拷贝。这个参数在训练规模较大时效果非常明显。
我自己的一个习惯是固定随机种子:
python复制import torch
import random
import numpy as np
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
因为数据增强里包含随机操作,如果不固定种子,每次跑实验数据分布都有细微差异,导致实验结果难以复现。固定种子之后,至少在同一个环境里能得到一致的结果。
4. models的妙用:迁移学习不再只是三行代码
4.1 预训练模型背后的那些“约定”
torchvision.models是目前复用视觉网络结构最直接的途径。你可以通过torchvision.models.resnet50(pretrained=True)加载在ImageNet上预训练过的权重。这个pretrained参数在较新版本中变成了weights参数,比如weights=ResNet50_Weights.IMAGENET1K_V1,但旧的布尔写法仍然兼容。实际上新接口更合理,因为你可以更精确地指定权重来源和版本。
迁移学习的核心逻辑是:预训练模型已经学习到了图像的基本纹理、边缘、形状等通用特征,你只需要在特定数据集上进行微调,让它把注意力集中在更任务相关的模式上。所以,加载模型时通常会默认冻结前几层权重,让它们保持已经学到的通用特征,只训练后面几层或新加的分类头。
我看到过很多人在微调时犯一个错误:直接把输入图片缩放成随机尺寸,然后喂给模型。这导致模型效果很差。原因在于,预训练模型要求输入尺寸与它训练时的输入尺寸一致。torchvision官方在模型文档里会写明每个模型期望的input_size,例如ResNet通常使用224x224。你按规定的尺寸做Resize和Crop后,模型才能发挥出预训练权重的作用。
4.2 更换分类头与冻结特征层
假设你有一个花分类数据集,总共10个类别。但预训练ResNet50最后一层全连接输出是1000维。你需要做的不只是替换最后一层,还得让模型的输出对齐类别数。最稳妥的写法:
python复制import torch.nn as nn
from torchvision import models
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10) # 换成10类分类头
如果你只想训练最后几层,可以把之前的所有参数requires_grad设为False:
python复制for param in model.parameters():
param.requires_grad = False
for param in model.fc.parameters():
param.requires_grad = True
这样反向传播时会跳过被冻结的层,大幅减少计算量和显存消耗。不过,如果新数据集与ImageNet差异较大(比如医学图像),只训练分类头往往不够,你可能需要解冻模型的后几层来做微调。解冻时要注意学习率,通常用较小的学习率(比如1e-4左右)进行微调,避免破坏预训练学到的特征。
4.3 不只是分类任务:torchvision也覆盖检测和分割
除了分类,torchvision.models还提供了检测模型(如Faster R-CNN、SSD)和分割模型(如DeepLabV3)。这些模型的API与分类模型不太一样,它们通常返回一个dict,包含boxes、labels、scores或masks等字段。
以目标检测为例,最基础的加载方式:
python复制from torchvision.models.detection import fasterrcnn_resnet50_fpn
model = fasterrcnn_resnet50_fpn(weights='DEFAULT')
使用这些模型时,数据预处理也和分类模型不同。检测模型通常期望输入是原始尺寸的图片张量,或者一个batch,但不需要归一化到ImageNet标准。你可以在官方文档中找到对应的预处理说明。我个人的体会是,检测和分割模型更适合在项目稳定后再深入研究,初学阶段先从分类入手,把数据管线理解透彻,再迁移到更复杂的任务会容易得多。
5. 一个完整的实战案例:从自定义数据集到微调ResNet
5.1 准备一个最小规模的数据集
理论讲再多,不如跑一个完整的流程。这里我准备一个实际做过的例子:用torchvision微调ResNet18完成一个“猫狗分类”的小项目。这个案例浓缩了前面讲到的所有核心步骤,你也可以换用自己的数据。
先准备一个简单的数据集。假设目录结构如下:
code复制data/cat_dog/
train/
cat/ # 若干猫图片
dog/ # 若干狗图片
val/
cat/
dog/
不需要太多图片,每类放20~50张就能跑通整个流程。真正要紧的是把数据加载逻辑搞对,模型训练反而是一道工序。
加载训练集和验证集时,分别应用不同的transform:
python复制from torchvision import datasets, transforms
train_transform = transforms.Compose([
transforms.Resize(256),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
val_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
train_dataset = datasets.ImageFolder('data/cat_dog/train', train_transform)
val_dataset = datasets.ImageFolder('data/cat_dog/val', val_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2)
这里RandomResizedCrop是训练时很有用的提升手段:它会随机裁剪一个区域再缩放到224x224,相当于让模型看到不同缩放比例和位置的物体。验证集不能这么做,否则无法公平衡量模型效果,所以用固定中心裁剪。
5.2 定义模型与训练循环
加载预训练模型,把分类头换成两个类别(猫和狗):
python复制from torchvision import models
import torch.optim as optim
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
num_features = model.fc.in_features
model.fc = torch.nn.Linear(num_features, 2)
这里把全部参数都设为可训练,因为对一个小型分类头来说,从头训练它需要梯度通过反向传播更新整个网络。但为了稳定迁移,我通常给分类头一个较大的学习率(比如1e-3),给特征提取部分一个较小的学习率(比如1e-4)。实现方式是把参数分成两组传给优化器:
python复制optimizer = optim.SGD([
{'params': model.features.parameters() if hasattr(model, 'features') else model.parameters(), 'lr': 1e-4},
{'params': model.fc.parameters(), 'lr': 1e-3},
], momentum=0.9)
但是注意,ResNet没有features这个属性,所以这里可以改成更灵活的方式——把model.layer1到model.layer4的参数视为特征提取部分,model.fc视为分类部分:
python复制params = []
for name, module in model.named_children():
if name == 'fc':
params.append({'params': module.parameters(), 'lr': 1e-3})
else:
params.append({'params': module.parameters(), 'lr': 1e-4})
optimizer = optim.SGD(params, momentum=0.9)
训练循环本身比较固定:
python复制criterion = torch.nn.CrossEntropyLoss()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
for epoch in range(30):
model.train()
running_loss = 0.0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
# 每个epoch后做验证
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in val_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Epoch {epoch+1}: train_loss={running_loss/len(train_dataset):.4f}, val_acc={correct/total:.4f}')
这里有个小经验:model.eval()和with torch.no_grad()是验证环节的标配,前者关闭Dropout和BatchNorm的统计更新,后者禁止梯度计算,能节省大量内存和计算时间。很多性能问题往往就是忘了调用model.eval(),导致验证时BatchNorm更新了running_mean,让结果飘忽不定。
5.3 训练结果分析与常见坑
这个最小的猫狗分类任务一般训练十几轮就能在验证集上有不错的表现(如果数据量很少,可能过拟合得很严重)。如果发现训练集准确率远高于验证集,说明过拟合了,此时可以增加数据增强的强度、引入Dropout或ReduceLROnPlateau等正则化手段。
常见坑之一是images=torch.Tensor形状不对。我在刚使用ResNet时,经常忘记将输入张量从[B,H,W,C]转为[B,C,H,W],导致维度不匹配。实际上torchvision的transforms里的ToTensor已经帮你做了转换,只要确保数据集返回的是[C,H,W]顺序即可。
另一个坑是模型没有调用.to(device),导致参数仍留在CPU,而数据去了GPU,训练直接报错。把模型和数据都挪到同一个设备上是最基本的规范。
6. 进阶技巧:性能优化、调试与个性化改造
6.1 加速训练:混合精度与梯度累积
如果你手里的GPU显存不大,混合精度是一个很有效的折中方案。PyTorch从1.6开始原生支持自动混合精度(AMP),你只需要修改训练循环中很少的代码:
python复制scaler = torch.cuda.amp.GradScaler()
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
AMP能利用GPU的Tensor Core,让训练速度提升20%到50%,同时显存占用也有降低。但要注意,某些算子不支持半精度,可能导致结果不准确。好消息是torchvision的模型结构基本都做了兼容处理,使用AMP一般不会出问题。
梯度累积是另一种绕过显存限制的方法。假设你的GPU最多只能放batch_size=16,但你想用batch_size=64的效果。可以在多个step上累积梯度,每隔4步才更新一次参数:
python复制accumulation_steps = 4
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(train_loader):
outputs = model(inputs)
loss = criterion(outputs, labels) / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
注意,每个sample的loss要除以累积步数,否则梯度会过大。这个方法在目标检测和分割任务中尤其常用,因为那些模型单个样本占显存就非常高。
6.2 用TensorBoard可视化工中间产物
调试深度学习模型时,可视化往往比打印日志更有效。torchvision配合TensorBoard可以很方便地查看输入图片、模型结构、Loss曲线和权重直方图。在训练循环里插入:
python复制from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('runs/exp1')
# 在训练前写一张网格图
grid = torchvision.utils.make_grid(inputs[:8])
writer.add_image('batch', grid, 0)
# 每个epoch记录loss和acc
writer.add_scalar('loss/train', running_loss / len(train_dataset), epoch)
writer.add_scalar('acc/val', correct / total, epoch)
torchvision.utils.make_grid是我非常喜欢的一个工具,它能把一个batch的多张图片拼接成网格图,保存下来后你可以直观检查数据增强的效果,看看图片有没有被过度裁剪、变色或者翻转。有时候模型训练效果差,先检查数据增强后的图片长什么样,往往能发现很多不合理的预处理参数。
6.3 保存与加载模型:不要只存状态字典
我见过太多人把自定义类和模型文件一起打包,结果加载时遇到反序列化问题。PyTorch的官方做法通常是保存state_dict而不是整个模型对象:
python复制torch.save(model.state_dict(), 'model.pth')
加载时,你需要先创建一个结构相同的模型,再加载state_dict:
python复制model = models.resnet18(weights=None)
model.fc = torch.nn.Linear(512, 2)
model.load_state_dict(torch.load('model.pth', map_location=device))
这里提一个来自热搜词“pytorch 2.6中weights_only参数”的注意点:从PyTorch 2.6开始,torch.load默认将weights_only=True,这是为了防止恶意pickle文件执行任意代码。如果加载的state_dict里只包含纯张量,是没问题的。但如果你曾保存过任何包含额外Python对象的state_dict(例如某个自定义类作为键),就可能需要显式设置weights_only=False,但这会带来安全风险。因此最好的实践是:不要在自己的state_dict里存奇怪的Python对象,只保存模型参数。
6.4 从源码中“捡漏”:自定义一个torchvision风格的工具
有时官方预置的功能不够用,比如我想在transforms里做正则化的同时加上mask处理。这种情况下可以模仿torchvision内部的写法,把一个可调用类封装成transform:
python复制class MyTransform:
def __init__(self, p=0.5):
self.p = p
def __call__(self, img, mask=None):
if torch.rand(1).item() < self.p:
# 自定义操作:同时处理img和mask
...
return img, mask
只要这个类的__call__方法接收并返回torchvision期望的数据类型,就能很好地嵌入到已有的Compose机制里。特别是在做语义分割时,对图片和mask需要做完全相同的随机翻转或裁剪,常规transforms.Compose无法同时操作两个目标,你必须自己写一个函数或类来处理。这个经验在我处理城市景观数据集时帮了大忙。
6.5 多卡训练时的torchvision模型同步
如果要用多卡训练,torchvision的模型需要配合torch.nn.DataParallel或DistributedDataParallel使用。简单场景下用DataParallel就够了:
python复制model = torch.nn.DataParallel(model)
model.to(device)
但要注意,DataParallel将batch拆分到多张卡,模型参数会有一份主副本,内存开销稍大。在更大规模的场景下,推荐使用DistributedDataParallel。无论用哪种,都需要注意BatchNorm层在数据分布发生变化时可能表现不稳定。如果你在多卡微调预训练模型,尤其要关注BatchNorm的同步均值/方差是否启用了同步(torch.nn.SyncBatchNorm.convert_sync_batchnorm),否则在多卡场景下每个卡的BatchNorm统计量不全局同步,可能让验证集准确率忽高忽低。
6.6 如何处理“官方文档与你手头环境不一致”的问题
最后想聊一个所有PyTorch使用者都会遇到的困扰:官方文档示例的API和你本地版本不一致。比如某些老版本使用pretrained=True,而新版本改成了weights参数。解决这个问题最简单的办法是查看当前torchvision源码。用print(torchvision.__file__)找到源码位置,直接grep关键参数。这个方法比上网搜更可靠,尤其当你的环境比较特殊时。
如果你需要修改源码包中的某个函数(比如想改变某个数据集的下载逻辑),可以继承原始类并覆盖方法,而不是直接改安装位置的文件。这样做能让你的代码在下次升级环境时不会崩溃,也方便维护。说到底,torchvision是一个“活”的库,它的接口版本在变,但底层逻辑相对稳定。掌握从源码中提炼信息的能力,比死记硬背某个函数签名长久得多。
最后再分享一个小技巧:在训练脚本顶部打印出torch、torchvision和Python的版本号,以及CUDA是否可用。有一次我调试到半夜,怎么都复现不了模型结果,最后发现是conda环境里不小心把torchvision升级到了另一个小版本,导致预训练权重下载地址都变了。从此以后,版本信息是每个实验记录的“第一行”,这个习惯让我少踩了无数个版本兼容的坑。
