★ FEATURED ARTICLE 前阵子有个朋友问我想用PyTorch做一个图像识别的小项目但他翻了一堆教程要么只跑通了单张图片的预测演示要么一上来就丢一个几百层的模型让人直接复制。我听完就觉得这类教程最大的问题是把“能跑”和“能落地”混为一谈了。所以这篇文章我不想复述官方文档而是把我实际做过的图像识别完整流程拆开讲从任务定义、数据准备、网络选型到训练调参、模型评估和部署前的最后一步。文章适合正在学PyTorch、打算拿它做图像识别的初学者也适合已经跑过一些Demo、但不确定下一步怎么走的中级使用者。我尽量把每一步背后的“为什么”也讲清楚而不是只给代码。1. 动手前先想清楚你的图像识别任务到底是什么1.1 先把“图像识别”翻译成计算机能解决的问题很多人上来就搜“PyTorch图像识别”这个词其实太宽泛了。图像识别在工程上至少会分化成图像分类、目标检测、图像分割、图像检索等好几个分支技术方案完全不一样。比如目标检测要输出“图中某个区域内有哪类目标”用的是边界框回归加分类图像分割则要像素级标记需要U-Net这类结构。如果你的需求只是“这张图是什么”那就对应单标签多分类问题这是入门最友好的一个方向。我建议动手之前先把任务参数写死类别数量、每类大概多少样本、输入图片的尺寸、图片里的目标是否足够显眼。另一个容易被忽略的点是数据来源。如果你只有几十张图几百张图就要考虑迁移学习如果每类有几千张甚至更多从头训练一个网络才说得通。把这些约束列清楚后面每一步选择都有据可依而不是凭感觉调参。1.2 为什么现在大多数学图像识别的人都选PyTorchPyTorch在图像识别领域流行核心原因是它的动态图机制。简单说你用代码写着写着就能随时打印中间张量的shape可以在中间层打断点看数据流这对调试神经网络来说太重要了。相比之下静态图模式相当于先整体编译再执行出了问题排错成本高很多。做图像识别的时候张量维度稍微搞错一点就报错能直接print出来看每一层输出排错速度完全不一样。另外PyTorch的生态对图像任务足够友好。torchvision集成了大量预训练模型、常用数据集和数据增强工具DataLoader配合Dataset能很舒服地处理图片数据。就算后续要部署到线上PyTorch也可以导出成ONNX格式给其他推理引擎消费。这些细节让我在实际项目里很少需要换框架。1.3 数据量、算力、精度目标决定你的方案选网络结构不能只看哪个模型刷榜刷得高要看你的资源和约束。我的经验可以分成三类情况供你参考。场景数据规模算力情况推荐方案快速验证每类几百张单块中端GPU预训练ResNet18 / MobileNet 只训练分类头常规项目每类几千张单块好点GPU预训练ResNet50微调最后几层 数据增强从零训练研究每类上万张多卡或长时间训练自己搭CNN或标准ResNet从头训这个决策逻辑很直接迁移学习是用别人的特征提取能力来弥补自己数据不足的短板相当于一个经验丰富的师傅先把“看边缘、看纹理、看形状”的基本功练好了你只需要教它识别你的特定类别。如果你数据量够大从头训练才可能超越预训练权重带来的上限。2. 环境与数据准备把底子打好2.1 环境安装里你大概率会遇到的版本问题我见过太多人卡在环境上不是PyTorch装不上而是装完之后运行时才发现版本对不上。最稳的流程是先创建一个独立的虚拟环境避免和系统Python或者其他项目互相污染。如果你用conda可以这样操作conda create -n classify python3.10 conda activate classify然后根据你是否使用GPU来安装对应版本。我的建议是无论你有没有独显都先在终端里跑下面这些命令确认一下python -c import torch; print(torch.__version__) python -c import torch; print(torch.cuda.is_available())如果第二行输出的是True那GPU环境就通了。很多人明明按教程安装但torch.cuda.is_available()依然是False最常见的原因是安装时选了CPU版本或者驱动版本太老。别急着重装系统先查驱动支持再去PyTorch官网选和你的CUDA能力匹配的安装命令这样最省时间。我踩过最冤的一个坑是conda自动匹配了CPU版torch结果代码跑了一天发现根本没用GPU浪费了大量时间。2.2 用文件夹结构组织图片比写几百行代码还管用数据准备阶段我最推荐的方式是把图片按类别放在文件夹里训练集、验证集、测试集分开目录管理。比如data/ train/ cat/ 图片... dog/ 图片... bird/ 图片... val/ cat/ dog/ bird/ test/ cat/ dog/ bird/这种结构的好处是torchvision.datasets.ImageFolder可以直接读取它会自动把子文件夹名称映射成类别标签连标签列表都不用自己写。如果数据是一个CSV或者数据库里存着路径和标签那就需要继承torch.utils.data.Dataset自定义读取逻辑但多数入门项目用ImageFolder就够了。from torchvision import datasets, transforms train_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_data datasets.ImageFolder(data/train, transformtrain_transform) print(train_data.classes) # [bird, cat, dog]注意Normalize里的均值和标准差。如果你用的是ImageNet预训练权重这三组数值几乎是固定的如果你从头训练自己的网络可以用全量训练数据估算均值和标准差效果更贴合你的数据分布。但很多新手最容易犯的错误就是对预训练模型用了和目标数据集完全不一致的归一化参数导致收敛困难。2.3 DataLoader参数逐个拆解batch_size、shuffle、num_workersDataLoader是连接数据集和训练循环之间的桥梁几个参数直接影响训练速度和效果。from torch.utils.data import DataLoader train_loader DataLoader( train_data, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue ) val_loader DataLoader( val_data, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue )batch_size每次迭代喂给GPU的图片数量。太大可能显存溢出太小会让梯度震荡训练不稳定。一般先在32或64里选OOM就减半验证集准确率迟迟上不去就试着加大。shuffle训练集必须置为True否则每一轮epoch内样本顺序不变模型可能学会记忆顺序而非特征验证集和测试集必须置为False因为评估结果不应该受打乱样本顺序影响。num_workers数据加载用了多少个进程。这个数值越高CPU读图和预处理越快但不是越大越好Windows上经常设置超过4会报错。pin_memory设为True会锁页内存把数据从CPU传输到GPU的耗时降低。这个参数在数据量大的时候提升明显代价是占用一部分内存。3. 网络结构设计先跑通一个简单的CNN再谈ResNet3.1 卷积网络是怎么“看”图像的如果你第一次接触卷积神经网络可以把它理解成一台滑动的扫描仪。卷积核就是在图片上从左到右、从上到下地滑动每次都只观察一个局部区域然后把局部区域的特征提取出来。相比全连接网络卷积层最大的优势是参数共享和局部连接。同一组卷积核在整个画幅上滑动意味着它只需要学习“某种纹理长什么样”而不是“某个位置上有这种纹理怎么处理”参数量大幅下降也更不容易过拟合。每个卷积层输出的特征图尺寸可以用这样一个基础公式估算输出边长 (输入边长 - 卷积核大小 2 * padding) / stride 1。比如输入224x224卷积核3x3padding为1stride为1输出就是224x224。池化层则是对特征图做降采样可以类比成把高清图片缩小成缩略图保留主要信息但丢掉冗余细节一定程度上也增强平移不变性。3.2 从零实现一个能跑的小型CNN我不建议初学者一上来就套大型网络先手写一个轻量CNN把这些概念串起来后面换ResNet才心里有底。下面这个结构是我常用的起步版输入3通道224x224图片输出3类import torch import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes3): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 224 - 112 nn.Conv2d(16, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 112 - 56 nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 56 - 28 ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(64, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x model SimpleCNN(num_classes3) dummy torch.randn(2, 3, 224, 224) out model(dummy) print(out.shape) # torch.Size([2, 3])我特意用了AdaptiveAvgPool2d它可以自适应地把任意大小的特征图池化成1x1这样你的模型对输入尺寸的宽容度会大很多换测试图片尺寸时不用改全连接层的维度。这个细节是我后面做跨分辨率预测时经常用到的。3.3 迁移学习把别人训练好的经验借过来自己从零训练一个CNN常常需要大量数据和很长时间实际项目中我更推荐直接在预训练模型基础上做迁移学习。torchvision里加载预训练模型很直接from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, num_classes) # 冻结大部分层只训练分类头 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True冻结策略有几种不同做法。最省资源和时间的是只训练最后一层分类头如果数据量中等可以尝试解冻后面几个block用较小的学习率训练如果你的数据已经足够多也可以不冻结任何层直接用预训练权重作为初始化全量微调。没有绝对优劣完全看你手里的数据量。我在实验中发现数据量只有每类一千张左右时只训分类头反而比全量微调更稳因为全量微调容易在早期就过拟合。4. 训练细节准确率提升的关键不在网络而在训练4.1 数据增强用更少的图片训练更强的模型很多人以为数据增强只是增加训练样本的手段实际上它还有很强的正则化作用。随机裁剪、水平翻转、颜色抖动这些变换相当于告诉模型“目标的类别不应该因为它在画面中的位置、方向、亮度变化而改变”。最常见的组合是训练集做RandomResizedCrop加RandomHorizontalFlip加ColorJitter验证集只做缩放和归一化。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]) ]) val_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]) ])注意验证集不要加随机增强否则每次测试结果都会因为随机操作而抖动没法稳定评估模型好坏。这是我最初最容易犯的一个原则性错误一度怀疑自己的模型不稳定。4.2 损失函数和优化器的选择逻辑图像分类最常用的损失函数是CrossEntropyLoss。它内部已经包含了Softmax的计算所以不需要再在网络最后加一个Softmax层。这个损失函数对多分类问题的效果和稳定性都很好我的建议是新手直接用它别折腾其他花式损失。优化器方面SGD和Adam是两大主流。Adam收敛快学习率可以设置得比较低尤其适合微调但SGD加上动量在训练后期往往有更强的泛化能力。我自己的习惯是跑通流程先用Adam学习率0.001验证思路没问题后再换SGD试试能否进一步提升准确率。还有一个容易被忽视的参数是weight_decay它本质上是给模型权重加了L2正则等于告诉模型“权重不要太极端”对控制过拟合有明显帮助。代码里就是一句optimizer torch.optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4)4.3 学习率训练中最常改的一个参数学习率设置太大会导致损失发散太小会让训练慢得让人怀疑人生。我踩过很多次学习率过高导致loss变成NaN的坑特别是用预训练模型解锁全量微调时整个网络学习率一致上调很容易震碎已经学好的权重。一个实用的经验是新加的分类头用相对大的学习率预训练特征提取部分用较小的学习率。optimizer torch.optim.Adam([ {params: model.fc.parameters(), lr: 1e-3}, {params: model.layer4.parameters(), lr: 1e-4}, ], weight_decay1e-4)学习率调度同样重要。我常用的是ReduceLROnPlateau也就是当验证集准确率不再上升时把学习率降为原来的十分之一。也可以用固定步长的StepLR最简单最省心。训练早期如果加一个小的“warmup”比如前面几个epoch把学习率从零逐渐线性升到目标值可以在使用较大初始学习率时不那么容易被撑坏。4.4 一个规范化的训练循环别把验证逻辑写在同一步里训练循环看起来简单但很多人把训练和验证混在一起处理搞到最后连模型是哪个epoch保存的都不知道。我的习惯是每轮epoch做两件事先在训练集上训练一遍再在验证集上完整评估一遍。两个函数分开写职责清晰。def train_one_epoch(model, loader, loss_fn, optimizer): model.train() total_loss, correct, total 0, 0, 0 for images, labels in loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss loss_fn(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() total labels.size(0) return total_loss / len(loader), correct / total def validate(model, loader, loss_fn): model.eval() total_loss, correct, total 0, 0, 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) loss loss_fn(outputs, labels) total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() total labels.size(0) return total_loss / len(loader), correct / total有两个细节我特别想说。第一model.train()和model.eval()必须在每个阶段开始前切换因为Dropout和BatchNorm在训练和推理时的行为完全不同。第二验证阶段一定用with torch.no_grad()包裹否则PyTorch会构建推理图显存瞬间被吃掉而且白白增加计算量。5. 模型评估与落地从测试集到真实项目5.1 不要只看准确率混淆矩阵和F1值能暴露更多问题很多初学者的唯一指标就是验证集准确率。这个指标在类别均衡时还凑合但遇到类别不均衡的数据集准确率会非常骗人。举个例子如果某个类占了90%的样本你的模型就算把所有图片都预测成这个类准确率也有90%实际一点用都没有。我建议每次训练完都看混淆矩阵把预测结果按行和列统计出来能直观看到哪些类别最容易互相混淆。然后计算每个类别的精确率、召回率、F1值尤其关注数量少的类别。用scikit-learn几行代码就能出来from sklearn.metrics import confusion_matrix, classification_report import numpy as np # preds和labels分别是所有预测结果和真实标签格式为List[int] cm confusion_matrix(labels, preds) print(cm) print(classification_report(labels, preds, target_names[cat, dog, bird]))如果两个类别在混淆矩阵里总是互相误判我一般会怀疑两类样本本身存在较大的视觉相似性或者标注质量有问题。这时候优先检查数据标注其次是增加对应类别的样本数量最后才考虑改网络结构。5.2 保存模型、加载模型和单张图片预测训练结束后保存模型有两种常见方式。只保存权重参数是最推荐的格式占空间小也方便后续换网络结构重新加载torch.save(model.state_dict(), model.pth)加载时要先构造一个同样结构的模型再把权重放进去model SimpleCNN(num_classes3) model.load_state_dict(torch.load(model.pth, map_locationcpu)) model.eval()推理阶段和训练阶段的写法完全不同。必须先切换到eval()模式再做一次和之前验证集相同的数据预处理然后调用模型from PIL import Image def predict_image(model, img_path, class_names): img Image.open(img_path).convert(RGB) img val_transform(img).unsqueeze(0) # 增加batch维度 model.eval() with torch.no_grad(): logits model(img) probs torch.softmax(logits, dim1) pred_idx logits.argmax(1).item() return class_names[pred_idx], probs[0][pred_idx].item() class_names train_data.classes pred, prob predict_image(model, test/cat_001.jpg, class_names) print(pred, prob)这里最容易翻车的点是忘了unsqueeze(0)。网络默认输入的shape是[batch, channel, height, width]单张图片读进来只有[channel, height, width]不加维度直接喂进去就报维度错误。另外一定要用softmax把输出转成概率分布否则你看到的就是一堆很难解释的数值。5.3 把模型接到实际项目里导成ONNX交给推理引擎如果只是做实验模型停在.pth文件就够了。真要把识别功能集成到一个Web服务或者App里我建议导出ONNX格式。ONNX是一种开放的模型交换格式很多推理引擎都能直接加载不依赖PyTorch环境本身。dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})这里设置dynamic_axes可以让输入的第一维变成动态的这样服务接口接收任意数量图片都行。导出之后再用ONNX Runtime跑一遍输入和输出都应该和PyTorch一致。能走到这一步一个图像识别模型基本就具备了被实际服务调用的能力。6. 复现与效果分析我踩过的坑和最终对比6.1 最容易踩的几个坑我按严重程度排了个序我梳理了自己在多个图像识别项目里反复踩过的坑也是经常在社区里看到新手们遇到的问题从严重到轻微排列如下。问题表现主要原因解决办法Normalize参数用错模型完全无法收敛loss居高不下和预训练模型不匹配核对来源模型标准归一化的均值和标准差训练时忘了model.train()使用带Dropout的网络时验证集很奇怪推理模式被用于训练每个阶段切换train/eval验证集被随机增强每次验证准确率波动大错误复用了训练预处理验证集只用ResizeCenterCropNormalize随机种子没固定每次训练结果完全不同初始化、数据采样顺序随机在代码顶部设置random.seed、np.random.seed、手动设置torch种子batch_size太大导致OOM训练中断显存溢出数据量和模型参数量超出显存减半batch_size或者混合精度其中最容易被忽视的是随机种子。如果你实验中发现连续两次训练同一个模型准确率差了1到2个百分点第一件事就是固定随机种子。下面这段代码我几乎每个项目都复制import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)6.2 同样一份数据下不同方案的效果差异有多大我给一个朋友做项目时用一份公开的图片分类数据集做了四组对比实验每类图片数量在800到1200张之间。跑出来的趋势很有参考价值方案训练轮数验证集准确率训练耗时从头训练SimpleCNN无增强4078%约20分钟从头训练SimpleCNN 数据增强4083%约25分钟预训练ResNet18 只训分类头1586%约15分钟预训练ResNet18 解冻layer4 增强2089%约20分钟这个对比的结论很清楚数据增强带来的提升比换网络结构更明显预训练模型哪怕只训分类头也比从头训练一个小网络准确率高出不少。而最后一种方案融合了增强和微调效果是最好的。实际解决业务问题时我通常沿用这个递进思路而不是一上来就追求最复杂的大模型。6.3 从头复现这套流程的操作顺序把上面内容收拢成一个可以照着做的清单我的建议顺序是定义分类任务确认类别数量和数据量。建立train/val/test目录核对每个类别的样本数和图片质量。固定随机种子搭建DataLoader和基础训练循环先用一个非常小的网络过一遍全流程。把网络换成预训练ResNet18先只训分类头确保能收敛。加入数据增强和学习率调度提升最终准确率。用混淆矩阵和F1检查每个类别的表现定位易混淆类别。保存模型导出ONNX接进实际应用。这套顺序的核心原则是先让流程转起来再逐步优化。我见过太多人一上来就在网络结构上花大量时间结果连基础的数据加载和训练循环都没跑通最后反而被一个维度错误卡了一整天。最后说一个我自己的习惯每改一次数据预处理、网络结构或者损失函数我都会先固定随机种子在小数据上跑几个epoch确认流程没问题再上全量数据。这个小习惯帮我过滤掉了大量“模型效果变差但原因却在于流程本身出bug”的虚假实验结论。图像识别这个方向看起来模型很炫真正决定项目成败的往往都是这些训练流程里的基本功。
阅读完成 · 觉得有帮助？
