最近有好几个朋友问我“用PyTorch搭一个分类网络到底该怎么从零开始做训练和测试”说实话这个问题我反复回答过很多次因为每个刚入门的人都会卡在同一个地方不是看不懂模型代码而是不清楚整个流程怎么串起来。我前前后后做过几个图像分类的小项目走过弯路也踩过坑今天干脆把这一整套东西整理出来。想认真学PyTorch分类任务的朋友这篇文章应该能帮你把数据、模型、训练、测试这条链路捋清楚照着做基本能跑通第一个项目。先说明一点我这里不会只贴代码还会讲清楚每个环节为什么这么做。因为分类网络看似简单但真正决定效果好坏、训练顺不顺的往往是那些文档里不会写的东西。比如数据预处理怎么配合预训练模型、训练时显存爆了怎么办、测试准确率和真实场景差多远这些都要心里有数。下面是正文。我会按照从数据到模型的顺序来写最后给出完整的代码示例和踩坑记录。所有代码我都用PyTorch 2.x版本跑过社区版、源码编译版都兼容但建议你先装好GPU版CPU版也能跑只是速度慢一些。1. 整体设计与网络选型思路1.1 为什么用PyTorch而不是其他框架现在做深度学习分类任务PyTorch基本是首选。它的动态图机制特别适合研究调试你可以随时print出中间张量的shape也可以在forward里写if语句这对刚接触神经网络的人非常友好。另一个优势是生态HuggingFace、TorchVision这些库都原生支持PyTorch你想要的预训练权重基本都有现成的。但选了PyTorch不等于万事大吉。我见过不少人上来就写一个很复杂的模型结果训练半天不收敛然后怪框架不好用。实际上大部分分类任务的难点不在框架而在数据处理和训练策略。所以这篇文章的重点也会放在这两个方向上。1.2 分类网络的核心组成一个标准的图像分类网络无论结构多复杂核心组成就三部分特征提取器、分类器、损失函数。特征提取器负责把原始图片变成高维特征比如一张224x224的彩色图经过卷积层和池化层后变成7x7x512的特征图分类器通常是一层或几层全连接把这个特征映射到类别得分上损失函数拿这些得分和真实标签对比算出差距来指导参数更新。初学者最容易犯的错误是拿着一个很深的网络从头训。如果没有大量数据深层网络很难收敛最后结果还不如一个浅层网络。所以我的建议是优先级从高到低排列的话先考虑数据质量和数量再用预训练模型微调最后才考虑自己设计网络结构。自己搭网络不是不行但对多数场景来说没必要重复造轮子。2. 数据准备与预处理2.1 数据目录结构与标准读取方式我习惯用一种约定俗成的目录结构来组织图片数据。每个类别一个文件夹文件夹名字就是类别标签。这样不需要额外写csv文件也不容易出错。data/ train/ cat/ 001.jpg dog/ 002.jpg val/ cat/ dog/如果数据量不大可以直接用torchvision.datasets.ImageFolder读取。它会自动把子文件夹名映射成数字标签返回一个Dataset对象。下面这个代码片段是训练集的标准读取方式from torchvision import datasets, transforms transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(rootdata/train, transformtransform) val_dataset datasets.ImageFolder(rootdata/val, transformtransform) print(train_dataset.classes) print(train_dataset.class_to_idx)这里有个非常关键的细节Normalize用的mean和std是ImageNet数据集的统计值。如果你要加载TorchVision里预训练好的模型就必须用这套归一化参数因为预训练模型是在同样的数据分布下训出来的。如果你用的是自己训练的网络归一化参数可以自己统计但前提是整个数据集统一标准。2.2 数据增强与训练验证集划分数据增强是提升分类模型泛化能力最直接的手段。常见的做法包括随机裁剪、随机水平翻转、随机旋转、颜色抖动。注意数据增强只应该用在训练集验证集和测试集只做缩放、中心裁剪和归一化。如果验证集也做了随机增强那你看到的指标会虚高因为你评估的并不是模型对真实图片的预测能力。train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])划分训练集和验证集时我建议不要在训练代码里用random_split而是把数据在文件系统层面分好。原因很简单你一次分好后面所有实验都共享同一份验证集结果才可以互相比较。如果每次跑实验都重新随机划分实验之间完全没有可比性。2.3 DataLoader的参数选择DataLoader是PyTorch中负责把Dataset包装成批量迭代器的工具。几个常用参数值得认真对待from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue)batch_size要根据显存来定。显卡显存8G的话224x224图片用32一版问题不大如果到了16G可以尝试64。num_workers一般设为CPU核心数或两倍核心数但也不是越大越好开太多有时反而会因为进程切换导致速度下降。pin_memoryTrue可以让数据传输到GPU更快因为锁页内存允许异步拷贝。我在这上面吃过亏。之前在一台老机器上把num_workers设成了8结果每个epoch刚开始就卡顿原因是机械硬盘IO跟不上多个worker同时抢读文件扇区寻道时间变长。后来把num_workers降到2速度反而上去了。所以数值不是死的要结合自己的硬件来看。3. 网络模型定义与关键参数3.1 从零搭建一个轻量CNN如果你想彻底弄懂分类网络内部发生了什么建议自己写一个很小的模型。下面这个模型有3个卷积块加一个全连接层适合小数据集比如10个类别的简单图像分类。import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes10): super(SimpleCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2, 2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2, 2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2, 2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * 28 * 28, 128), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x注意最后的全连接层没有加Softmax。很多初学者会在这里加然后在损失函数里又用CrossEntropyLoss导致训练出错。实际上PyTorch的CrossEntropyLoss内部已经包含了Softmax计算你在网络输出层直接给原始logits就行。如果加了Softmax反而会破坏数值稳定性训练时有可能梯度消失。3.2 使用预训练模型微调如果数据量中等比如每类500张图以上强烈建议用TorchVision里预训练好的ResNet18或ResNet50。加载方式非常简单from torchvision import models model models.resnet18(pretrainedTrue) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes)这样做的好处是模型底层的卷积核已经学会了提取通用图像特征比如边缘、纹理、形状。你要做的只是让最后一层全连接学会把这些特征映射到你的类别上。微调的时候有两种策略只训练最后一层冻结前面的卷积层。适合数据量很少、且你的图片分布和ImageNet比较接近的场景。全模型微调使用较小的学习率。适合数据量中等或较多的场景。我一般偏向后者因为只训练最后一层可能学到的特征和你的任务不匹配。但全模型微调要小心过拟合数据增强必须跟上。3.3 损失函数与优化器怎么选分类任务最常用的损失函数就是交叉熵损失PyTorch对应CrossEntropyLoss。多分类基本用它没有太多悬念。优化器的选择就比较有意思了。SGD加动量一直是经典选择学习率调到0.01左右配合余弦退火或者StepLR很多任务都能收敛得很好。Adam适合快速出结果但有时候会收敛到比较平的极小值泛化能力稍差。我自己的习惯是如果训练时间充裕用SGD如果只是先跑通流程用Adam。criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30)weight_decay是非常重要但容易被忽略的参数。它本质上是L2正则化可以压制大权重减轻过拟合。对于图像分类我一般设1e-4或5e-5。如果你发现训练集准确率很高但验证集上不去先把weight_decay调大试试。4. 训练流程实现4.1 训练循环的完整代码下面这个训练函数是我项目里一直在用的逻辑非常清晰模型先进入训练模式遍历DataLoader把图片和标签搬到GPU前向计算损失反向传播更新梯度最后统计这个epoch的平均损失和准确率。不要忘了每个epoch开始前要把模型设置成model.train()每个epoch结束后设置成model.eval()。import torch def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in train_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total acc correct / total return avg_loss, acc这里有两个容易出问题的点。一是optimizer.zero_grad()必须放在loss.backward()之前否则梯度会累加到上一轮。二是batch里的loss要乘上batch size再累加最后除总数这样就算最后一个batch不满也能得到正确的平均损失。有些人直接用loss.item()累加最后除以len(dataloader)我建议不要这样因为最后一个batch可能小于batch_size结果会有偏差。4.2 验证函数的写法验证函数和训练函数类似但有几个关键区别不需要计算梯度不需要更新参数需要使用torch.no_grad()包裹同时一定要model.eval()。因为在eval模式下Dropout会被关闭BatchNorm也会使用全局统计量而不是当前batch的统计量。def validate(model, val_loader, criterion, device): model.eval() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total acc correct / total return avg_loss, acc我记得最初写验证函数的时候忘了加model.eval()结果验证集准确率波动特别大有时候甚至比训练集还高。原因就是Dropout没有关闭模型每次前向都是随机的。这个问题看起来不大但足以让整个实验结论失真。4.3 训练主循环与模型保存训练主循环就是迭代多个epoch每个epoch里跑一遍训练和验证然后记录日志。我通常会把效果最好的模型单独保存下来而不是只在最后保存一次。因为训练后期可能过拟合最后那个epoch往往不是验证集上表现最好的。import os best_acc 0.0 num_epochs 30 device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) for epoch in range(num_epochs): train_loss, train_acc train_one_epoch( model, train_loader, criterion, optimizer, device ) val_loss, val_acc validate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch1}/{num_epochs} f Train Loss: {train_loss:.4f} Train Acc: {train_acc:.4f} f Val Loss: {val_loss:.4f} Val Acc: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print( - Saved best model)保存模型时用state_dict()而不是直接保存整个model对象。state_dict只包含参数不包含模型结构所以加载时你需要先实例化一个一样的模型结构再load_state_dict。直接保存整个model的好处是方便但坏处是耦合了代码结构以后换环境、换GPU可能就加载不了。我推荐只保存state_dict。4.4 学习率调整策略学习率调度是训练成败的分水岭。常见的有StepLR、ReduceLROnPlateau和CosineAnnealingLR。StepLR简单粗暴每隔固定步数把学习率乘以一个因子ReduceLROnPlateau在监控指标停滞时自动降低学习率CosineAnnealingLR则让学习率按照余弦曲线从大到小衰减在训练后期放缓更新步伐有助于收敛到更好的局部最优。我比较常用CosineAnnealingLR因为它不需要太关注在哪个epoch停止整体趋势是越来越小。但注意如果epoch数设得太短学习率可能还没降到足够低就到了末尾效果反而不如StepLR稳定。自己动手做实验的话可以多试几种配合TensorBoard看loss曲线选择最适合你数据量的那一种。5. 测试评估与结果分析5.1 测试集上的模型加载与评估测试阶段和验证阶段本质上一样但测试集应该是模型完全没见过、只在最后评估时用一遍的数据。加载模型时要注意路径和结构一致代码如下model models.resnet18(pretrainedFalse) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model model.to(device) test_loss, test_acc validate(model, test_loader, criterion, device) print(fTest Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f})map_locationdevice非常重要。如果你的模型在GPU上训练测试机器只有CPU不加这个参数会报模型参数都在GPU上的错误。反过来也一样在CPU上训练的模型要加载到GPU也要指定map_location。5.2 用混淆矩阵分析错误类型只看一个整体准确率其实会掩盖很多问题。比如你有10个类别其中某个类别准确率特别低它的样本全被分到了另一个类里这时候准确率可能还挺高但实际模型是偏的。所以建议测试完以后画出混淆矩阵看哪些类别互相混淆。import numpy as np from sklearn.metrics import confusion_matrix model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in test_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) print(cm)有了混淆矩阵你就能很直观地看到哪些类别语义相近比如猫和狗、卡车和汽车。如果模型老是把A类识别成B类可以考虑增加A类样本的拍摄角度、在数据增强里加入色彩扰动或者给A类样本适当加权。5.3 单张图片推理示例训练完模型后真正部署时需要写一个单张图片推理的函数。这里有几个容易踩的坑图片要解码成RGB要缩放成训练时相同的尺寸要使用相同的归一化参数要加一个batch维度。下面是一个可直接用的示例from PIL import Image import torchvision.transforms as transforms def predict_image(image_path, model, class_names, device): transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) image transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): outputs model(image) _, predicted torch.max(outputs, 1) probs torch.softmax(outputs, dim1) class_idx predicted.item() confidence probs[0, class_idx].item() return class_names[class_idx], confidenceclass_names可以从之前ImageFolder的classes属性里拿顺序要和训练时完全一致。如果顺序不一致预测结果就会全部错位这种错误特别隐蔽我见过有人排查了一整天最后发现是测试时自己手写了一个类别列表顺序搞错了。6. 常见问题与排查技巧实录6.1 训练loss不下降或者直接NAN先说说NAN的情况。最常见的原因是学习率设置得太大。如果你用的是Adam学习率默认是0.001但有些新版本或插件会自动调整一旦学习率超过模型承受范围loss就会在某个step变成NAN。解决办法是把学习率调到原来的1/10重新试。还有一种可能是数据里有损坏的图片或无穷大的像素值可以在Dataset的__getitem__里加入异常处理打印出错路径。如果loss不下降但不是NAN那要先看数据本身。类别是否严重不平衡像素值有没有归一化标签是否从0开始连续编号另外连续使用BatchNorm时如果batch_size太小比如只有1或2也会导致训练不稳定因为统计量方差太大。这种情况下要么增大batch_size要么干脆不用BatchNorm。6.2 过拟合的识别与处理训练集准确率超过98%验证集只有80%左右这就是典型的过拟合。缓解手段按优先级排列用更强的数据增强、增加Dropout、增大weight_decay、减小模型复杂度、提前停止训练。如果这些都不管用就该去收集更多数据了数据增强只能算是“无中生有”的替代方案。有一个小技巧可以观察训练和验证的loss曲线。如果验证loss先下降后上升准确率却还在微涨此时模型已经开始记住训练集的噪声应该立刻停止训练以验证loss最低点的模型为准。这就是为什么每一步验证并保存最优模型极其重要而不是到全部训练完再保存。6.3 显存不足的应对方案显存不足有两种报错CUDA out of memory和OOM。最简单的方法是减小batch_size这是立竿见影的办法。如果batch_size已经小到1还是不够可以检查一下输入图片的分辨率是不是太大了。另外记得在验证和测试的时候用torch.no_grad()否则会为不需要反传的图保留中间梯度。还有一个容易被忽视的点在DataLoader里开了many workers但每个worker都会持有数据副本显存占用会额外增加。如果开了pin_memory也会有一些显存开销。我建议先关掉这些选项排查再逐个打开。6.4 分类结果在某些类别上特别差这种情况建议画一个每个类别的单独准确率直方图。你会发现差的那几个类别往往样本数量少或者图片风格和别的类很接近。不要急着改模型可以先做类别加权在损失函数中给少数类更大的权重。PyTorch的CrossEntropyLoss支持weight参数可以直接传入一个tensorweights torch.tensor([1.0, 1.0, 5.0, ...]).to(device) criterion nn.CrossEntropyLoss(weightweights)权重的反比设计通常是权重 总样本数 / (类别数 * 该类样本数)。但这也不是固定公式需要根据实际效果调整。如果加了权重后总体准确率下降说明少数类占比太小强行学习它们的代价是牺牲多数类的性能这时候更好的办法是收集更多少数类样本或者使用更激进的数据增强。6.5 训练速度很慢怎么办如果GPU利用率不高问题通常出在数据读取上。先检查num_workers试着调高一些再检查是否是CPU瓶颈可以观察任务管理器。如果数据需要每次都做大量随机裁剪和颜色抖动CPU成本确实很高可以考虑用CUDA加速的预处理。但多数情况下把图片在进入Dataset前缩放到较小尺寸比如256x256再在训练时随机裁剪到224x224已经能显著减少IO和CPU压力。还有一点别在训练循环里频繁print大量信息每轮只print一行就好。如果每次print都输出很多内容GPU会等你打印完才开始下一个batch速度就拖慢了。我习惯让训练函数返回累计的指标最后统一打印效率会好很多。6.6 关于可复现性的建议想让实验结果可以复现需要固定三样东西PyTorch随机种子、Python随机种子、NumPy随机种子。另外还要对DataLoader设置一个固定的worker随机种子不过PyTorch在2.x版本里会自动处理好大部分。下面这段代码可以放到训练脚本开头import random import numpy as np import torch def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False set_seed(42)注意cudnn.deterministicTrue会在一定程度上降低速度因为卷积算法失去了自动调优的空间。实验阶段为了可复现可以开但真正追求训练速度时建议关掉。7. 一些我踩过坑之后总结的心得做分类网络这件事模型结构只占很小一部分数据、训练策略、评估方式才是决定成败的关键。我见过很多人把大量时间花在调模型层数上结果数据预处理随便写写最后效果一塌糊涂。我的建议是先把整套流程跑通再一步步优化每一步的改动都要同时看到训练集和验证集的变化不要只看最终测试准确率。另外日志记录要趁早做。每跑一个实验把超参数、数据增强方式、模型结构、最终准确率都记下来。后面你会遇到这种情况一个配置看上去效果不错但你说不清楚它到底加了多少程度的数据增强想要复现就得重新猜。所以强烈建议用类似TensorBoard或简单的csv表格记录每个实验的关键信息这一点真的能为你节省大量时间。如果你是想用分类网络做实际项目还有一点要注意真实场景下的图片往往和数据集里分布不一致摄像头角度、光照、背景都会影响结果。所以在测试集上准确率再高也不代表上线后就稳。建议在部署前收集一小批真实场景图片做一次额外的冒烟测试看看模型能不能泛化过去。这篇文章里给的代码是我在多个项目里反复使用并验证过的基本可以直接搬过去用。但别忘了调参尤其是学习率和batch_size这两个参数直接影响训练效果需要根据自己的数据和显存重新实验。希望这篇内容能帮你把PyTorch分类网络从搭建到上线的过程走通少走一些我当年走过的弯路。 
长河君 · 2026/10/16 5:07:55 · 阅读 1.2k PyTorch图像分类实战：从数据预处理到模型训练测试全流程
最近有好几个朋友问我“用PyTorch搭一个分类网络到底该怎么从零开始做训练和测试”说实话这个问题我反复回答过很多次因为每个刚入门的人都会卡在同一个地方不是看不懂模型代码而是不清楚整个流程怎么串起来。我前前后后做过几个图像分类的小项目走过弯路也踩过坑今天干脆把这一整套东西整理出来。想认真学PyTorch分类任务的朋友这篇文章应该能帮你把数据、模型、训练、测试这条链路捋清楚照着做基本能跑通第一个项目。先说明一点我这里不会只贴代码还会讲清楚每个环节为什么这么做。因为分类网络看似简单但真正决定效果好坏、训练顺不顺的往往是那些文档里不会写的东西。比如数据预处理怎么配合预训练模型、训练时显存爆了怎么办、测试准确率和真实场景差多远这些都要心里有数。下面是正文。我会按照从数据到模型的顺序来写最后给出完整的代码示例和踩坑记录。所有代码我都用PyTorch 2.x版本跑过社区版、源码编译版都兼容但建议你先装好GPU版CPU版也能跑只是速度慢一些。1. 整体设计与网络选型思路1.1 为什么用PyTorch而不是其他框架现在做深度学习分类任务PyTorch基本是首选。它的动态图机制特别适合研究调试你可以随时print出中间张量的shape也可以在forward里写if语句这对刚接触神经网络的人非常友好。另一个优势是生态HuggingFace、TorchVision这些库都原生支持PyTorch你想要的预训练权重基本都有现成的。但选了PyTorch不等于万事大吉。我见过不少人上来就写一个很复杂的模型结果训练半天不收敛然后怪框架不好用。实际上大部分分类任务的难点不在框架而在数据处理和训练策略。所以这篇文章的重点也会放在这两个方向上。1.2 分类网络的核心组成一个标准的图像分类网络无论结构多复杂核心组成就三部分特征提取器、分类器、损失函数。特征提取器负责把原始图片变成高维特征比如一张224x224的彩色图经过卷积层和池化层后变成7x7x512的特征图分类器通常是一层或几层全连接把这个特征映射到类别得分上损失函数拿这些得分和真实标签对比算出差距来指导参数更新。初学者最容易犯的错误是拿着一个很深的网络从头训。如果没有大量数据深层网络很难收敛最后结果还不如一个浅层网络。所以我的建议是优先级从高到低排列的话先考虑数据质量和数量再用预训练模型微调最后才考虑自己设计网络结构。自己搭网络不是不行但对多数场景来说没必要重复造轮子。2. 数据准备与预处理2.1 数据目录结构与标准读取方式我习惯用一种约定俗成的目录结构来组织图片数据。每个类别一个文件夹文件夹名字就是类别标签。这样不需要额外写csv文件也不容易出错。data/ train/ cat/ 001.jpg dog/ 002.jpg val/ cat/ dog/如果数据量不大可以直接用torchvision.datasets.ImageFolder读取。它会自动把子文件夹名映射成数字标签返回一个Dataset对象。下面这个代码片段是训练集的标准读取方式from torchvision import datasets, transforms transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(rootdata/train, transformtransform) val_dataset datasets.ImageFolder(rootdata/val, transformtransform) print(train_dataset.classes) print(train_dataset.class_to_idx)这里有个非常关键的细节Normalize用的mean和std是ImageNet数据集的统计值。如果你要加载TorchVision里预训练好的模型就必须用这套归一化参数因为预训练模型是在同样的数据分布下训出来的。如果你用的是自己训练的网络归一化参数可以自己统计但前提是整个数据集统一标准。2.2 数据增强与训练验证集划分数据增强是提升分类模型泛化能力最直接的手段。常见的做法包括随机裁剪、随机水平翻转、随机旋转、颜色抖动。注意数据增强只应该用在训练集验证集和测试集只做缩放、中心裁剪和归一化。如果验证集也做了随机增强那你看到的指标会虚高因为你评估的并不是模型对真实图片的预测能力。train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])划分训练集和验证集时我建议不要在训练代码里用random_split而是把数据在文件系统层面分好。原因很简单你一次分好后面所有实验都共享同一份验证集结果才可以互相比较。如果每次跑实验都重新随机划分实验之间完全没有可比性。2.3 DataLoader的参数选择DataLoader是PyTorch中负责把Dataset包装成批量迭代器的工具。几个常用参数值得认真对待from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue)batch_size要根据显存来定。显卡显存8G的话224x224图片用32一版问题不大如果到了16G可以尝试64。num_workers一般设为CPU核心数或两倍核心数但也不是越大越好开太多有时反而会因为进程切换导致速度下降。pin_memoryTrue可以让数据传输到GPU更快因为锁页内存允许异步拷贝。我在这上面吃过亏。之前在一台老机器上把num_workers设成了8结果每个epoch刚开始就卡顿原因是机械硬盘IO跟不上多个worker同时抢读文件扇区寻道时间变长。后来把num_workers降到2速度反而上去了。所以数值不是死的要结合自己的硬件来看。3. 网络模型定义与关键参数3.1 从零搭建一个轻量CNN如果你想彻底弄懂分类网络内部发生了什么建议自己写一个很小的模型。下面这个模型有3个卷积块加一个全连接层适合小数据集比如10个类别的简单图像分类。import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes10): super(SimpleCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2, 2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2, 2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2, 2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * 28 * 28, 128), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x注意最后的全连接层没有加Softmax。很多初学者会在这里加然后在损失函数里又用CrossEntropyLoss导致训练出错。实际上PyTorch的CrossEntropyLoss内部已经包含了Softmax计算你在网络输出层直接给原始logits就行。如果加了Softmax反而会破坏数值稳定性训练时有可能梯度消失。3.2 使用预训练模型微调如果数据量中等比如每类500张图以上强烈建议用TorchVision里预训练好的ResNet18或ResNet50。加载方式非常简单from torchvision import models model models.resnet18(pretrainedTrue) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes)这样做的好处是模型底层的卷积核已经学会了提取通用图像特征比如边缘、纹理、形状。你要做的只是让最后一层全连接学会把这些特征映射到你的类别上。微调的时候有两种策略只训练最后一层冻结前面的卷积层。适合数据量很少、且你的图片分布和ImageNet比较接近的场景。全模型微调使用较小的学习率。适合数据量中等或较多的场景。我一般偏向后者因为只训练最后一层可能学到的特征和你的任务不匹配。但全模型微调要小心过拟合数据增强必须跟上。3.3 损失函数与优化器怎么选分类任务最常用的损失函数就是交叉熵损失PyTorch对应CrossEntropyLoss。多分类基本用它没有太多悬念。优化器的选择就比较有意思了。SGD加动量一直是经典选择学习率调到0.01左右配合余弦退火或者StepLR很多任务都能收敛得很好。Adam适合快速出结果但有时候会收敛到比较平的极小值泛化能力稍差。我自己的习惯是如果训练时间充裕用SGD如果只是先跑通流程用Adam。criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30)weight_decay是非常重要但容易被忽略的参数。它本质上是L2正则化可以压制大权重减轻过拟合。对于图像分类我一般设1e-4或5e-5。如果你发现训练集准确率很高但验证集上不去先把weight_decay调大试试。4. 训练流程实现4.1 训练循环的完整代码下面这个训练函数是我项目里一直在用的逻辑非常清晰模型先进入训练模式遍历DataLoader把图片和标签搬到GPU前向计算损失反向传播更新梯度最后统计这个epoch的平均损失和准确率。不要忘了每个epoch开始前要把模型设置成model.train()每个epoch结束后设置成model.eval()。import torch def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in train_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total acc correct / total return avg_loss, acc这里有两个容易出问题的点。一是optimizer.zero_grad()必须放在loss.backward()之前否则梯度会累加到上一轮。二是batch里的loss要乘上batch size再累加最后除总数这样就算最后一个batch不满也能得到正确的平均损失。有些人直接用loss.item()累加最后除以len(dataloader)我建议不要这样因为最后一个batch可能小于batch_size结果会有偏差。4.2 验证函数的写法验证函数和训练函数类似但有几个关键区别不需要计算梯度不需要更新参数需要使用torch.no_grad()包裹同时一定要model.eval()。因为在eval模式下Dropout会被关闭BatchNorm也会使用全局统计量而不是当前batch的统计量。def validate(model, val_loader, criterion, device): model.eval() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total acc correct / total return avg_loss, acc我记得最初写验证函数的时候忘了加model.eval()结果验证集准确率波动特别大有时候甚至比训练集还高。原因就是Dropout没有关闭模型每次前向都是随机的。这个问题看起来不大但足以让整个实验结论失真。4.3 训练主循环与模型保存训练主循环就是迭代多个epoch每个epoch里跑一遍训练和验证然后记录日志。我通常会把效果最好的模型单独保存下来而不是只在最后保存一次。因为训练后期可能过拟合最后那个epoch往往不是验证集上表现最好的。import os best_acc 0.0 num_epochs 30 device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) for epoch in range(num_epochs): train_loss, train_acc train_one_epoch( model, train_loader, criterion, optimizer, device ) val_loss, val_acc validate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch1}/{num_epochs} f Train Loss: {train_loss:.4f} Train Acc: {train_acc:.4f} f Val Loss: {val_loss:.4f} Val Acc: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print( - Saved best model)保存模型时用state_dict()而不是直接保存整个model对象。state_dict只包含参数不包含模型结构所以加载时你需要先实例化一个一样的模型结构再load_state_dict。直接保存整个model的好处是方便但坏处是耦合了代码结构以后换环境、换GPU可能就加载不了。我推荐只保存state_dict。4.4 学习率调整策略学习率调度是训练成败的分水岭。常见的有StepLR、ReduceLROnPlateau和CosineAnnealingLR。StepLR简单粗暴每隔固定步数把学习率乘以一个因子ReduceLROnPlateau在监控指标停滞时自动降低学习率CosineAnnealingLR则让学习率按照余弦曲线从大到小衰减在训练后期放缓更新步伐有助于收敛到更好的局部最优。我比较常用CosineAnnealingLR因为它不需要太关注在哪个epoch停止整体趋势是越来越小。但注意如果epoch数设得太短学习率可能还没降到足够低就到了末尾效果反而不如StepLR稳定。自己动手做实验的话可以多试几种配合TensorBoard看loss曲线选择最适合你数据量的那一种。5. 测试评估与结果分析5.1 测试集上的模型加载与评估测试阶段和验证阶段本质上一样但测试集应该是模型完全没见过、只在最后评估时用一遍的数据。加载模型时要注意路径和结构一致代码如下model models.resnet18(pretrainedFalse) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, num_classes) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model model.to(device) test_loss, test_acc validate(model, test_loader, criterion, device) print(fTest Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f})map_locationdevice非常重要。如果你的模型在GPU上训练测试机器只有CPU不加这个参数会报模型参数都在GPU上的错误。反过来也一样在CPU上训练的模型要加载到GPU也要指定map_location。5.2 用混淆矩阵分析错误类型只看一个整体准确率其实会掩盖很多问题。比如你有10个类别其中某个类别准确率特别低它的样本全被分到了另一个类里这时候准确率可能还挺高但实际模型是偏的。所以建议测试完以后画出混淆矩阵看哪些类别互相混淆。import numpy as np from sklearn.metrics import confusion_matrix model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in test_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) print(cm)有了混淆矩阵你就能很直观地看到哪些类别语义相近比如猫和狗、卡车和汽车。如果模型老是把A类识别成B类可以考虑增加A类样本的拍摄角度、在数据增强里加入色彩扰动或者给A类样本适当加权。5.3 单张图片推理示例训练完模型后真正部署时需要写一个单张图片推理的函数。这里有几个容易踩的坑图片要解码成RGB要缩放成训练时相同的尺寸要使用相同的归一化参数要加一个batch维度。下面是一个可直接用的示例from PIL import Image import torchvision.transforms as transforms def predict_image(image_path, model, class_names, device): transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) image transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): outputs model(image) _, predicted torch.max(outputs, 1) probs torch.softmax(outputs, dim1) class_idx predicted.item() confidence probs[0, class_idx].item() return class_names[class_idx], confidenceclass_names可以从之前ImageFolder的classes属性里拿顺序要和训练时完全一致。如果顺序不一致预测结果就会全部错位这种错误特别隐蔽我见过有人排查了一整天最后发现是测试时自己手写了一个类别列表顺序搞错了。6. 常见问题与排查技巧实录6.1 训练loss不下降或者直接NAN先说说NAN的情况。最常见的原因是学习率设置得太大。如果你用的是Adam学习率默认是0.001但有些新版本或插件会自动调整一旦学习率超过模型承受范围loss就会在某个step变成NAN。解决办法是把学习率调到原来的1/10重新试。还有一种可能是数据里有损坏的图片或无穷大的像素值可以在Dataset的__getitem__里加入异常处理打印出错路径。如果loss不下降但不是NAN那要先看数据本身。类别是否严重不平衡像素值有没有归一化标签是否从0开始连续编号另外连续使用BatchNorm时如果batch_size太小比如只有1或2也会导致训练不稳定因为统计量方差太大。这种情况下要么增大batch_size要么干脆不用BatchNorm。6.2 过拟合的识别与处理训练集准确率超过98%验证集只有80%左右这就是典型的过拟合。缓解手段按优先级排列用更强的数据增强、增加Dropout、增大weight_decay、减小模型复杂度、提前停止训练。如果这些都不管用就该去收集更多数据了数据增强只能算是“无中生有”的替代方案。有一个小技巧可以观察训练和验证的loss曲线。如果验证loss先下降后上升准确率却还在微涨此时模型已经开始记住训练集的噪声应该立刻停止训练以验证loss最低点的模型为准。这就是为什么每一步验证并保存最优模型极其重要而不是到全部训练完再保存。6.3 显存不足的应对方案显存不足有两种报错CUDA out of memory和OOM。最简单的方法是减小batch_size这是立竿见影的办法。如果batch_size已经小到1还是不够可以检查一下输入图片的分辨率是不是太大了。另外记得在验证和测试的时候用torch.no_grad()否则会为不需要反传的图保留中间梯度。还有一个容易被忽视的点在DataLoader里开了many workers但每个worker都会持有数据副本显存占用会额外增加。如果开了pin_memory也会有一些显存开销。我建议先关掉这些选项排查再逐个打开。6.4 分类结果在某些类别上特别差这种情况建议画一个每个类别的单独准确率直方图。你会发现差的那几个类别往往样本数量少或者图片风格和别的类很接近。不要急着改模型可以先做类别加权在损失函数中给少数类更大的权重。PyTorch的CrossEntropyLoss支持weight参数可以直接传入一个tensorweights torch.tensor([1.0, 1.0, 5.0, ...]).to(device) criterion nn.CrossEntropyLoss(weightweights)权重的反比设计通常是权重 总样本数 / (类别数 * 该类样本数)。但这也不是固定公式需要根据实际效果调整。如果加了权重后总体准确率下降说明少数类占比太小强行学习它们的代价是牺牲多数类的性能这时候更好的办法是收集更多少数类样本或者使用更激进的数据增强。6.5 训练速度很慢怎么办如果GPU利用率不高问题通常出在数据读取上。先检查num_workers试着调高一些再检查是否是CPU瓶颈可以观察任务管理器。如果数据需要每次都做大量随机裁剪和颜色抖动CPU成本确实很高可以考虑用CUDA加速的预处理。但多数情况下把图片在进入Dataset前缩放到较小尺寸比如256x256再在训练时随机裁剪到224x224已经能显著减少IO和CPU压力。还有一点别在训练循环里频繁print大量信息每轮只print一行就好。如果每次print都输出很多内容GPU会等你打印完才开始下一个batch速度就拖慢了。我习惯让训练函数返回累计的指标最后统一打印效率会好很多。6.6 关于可复现性的建议想让实验结果可以复现需要固定三样东西PyTorch随机种子、Python随机种子、NumPy随机种子。另外还要对DataLoader设置一个固定的worker随机种子不过PyTorch在2.x版本里会自动处理好大部分。下面这段代码可以放到训练脚本开头import random import numpy as np import torch def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False set_seed(42)注意cudnn.deterministicTrue会在一定程度上降低速度因为卷积算法失去了自动调优的空间。实验阶段为了可复现可以开但真正追求训练速度时建议关掉。7. 一些我踩过坑之后总结的心得做分类网络这件事模型结构只占很小一部分数据、训练策略、评估方式才是决定成败的关键。我见过很多人把大量时间花在调模型层数上结果数据预处理随便写写最后效果一塌糊涂。我的建议是先把整套流程跑通再一步步优化每一步的改动都要同时看到训练集和验证集的变化不要只看最终测试准确率。另外日志记录要趁早做。每跑一个实验把超参数、数据增强方式、模型结构、最终准确率都记下来。后面你会遇到这种情况一个配置看上去效果不错但你说不清楚它到底加了多少程度的数据增强想要复现就得重新猜。所以强烈建议用类似TensorBoard或简单的csv表格记录每个实验的关键信息这一点真的能为你节省大量时间。如果你是想用分类网络做实际项目还有一点要注意真实场景下的图片往往和数据集里分布不一致摄像头角度、光照、背景都会影响结果。所以在测试集上准确率再高也不代表上线后就稳。建议在部署前收集一小批真实场景图片做一次额外的冒烟测试看看模型能不能泛化过去。这篇文章里给的代码是我在多个项目里反复使用并验证过的基本可以直接搬过去用。但别忘了调参尤其是学习率和batch_size这两个参数直接影响训练效果需要根据自己的数据和显存重新实验。希望这篇内容能帮你把PyTorch分类网络从搭建到上线的过程走通少走一些我当年走过的弯路。