这期填一个之前挖的坑用 Pytorch 搭一个图像分类器专门识别照片里的运动鞋是哪家品牌的。Pytorch 这套框架我用了很久从早期 0.4 版本一路跟到 2.x不管是做实验还是快速验证想法它都算得上顺手。这次项目选择了三分类任务输入一张运动鞋照片模型输出品牌标签其实换到别的商品类别也完全可以平移比如识别球衣、背包甚至宠物品种原理完全一致。先交代一下这个项目适合谁来参考如果你刚学完 Pytorch 基础框架想找一个“能跑通、看得懂、能改着玩”的图像分类练手项目这一篇可以拿来直接用如果你已经写过不少分类模型也可以看看我在数据处理和调参环节踩过的坑。整个项目完整走了一遍数据准备、模型搭建、训练评估和问题排查算是非常标准的流程。1. 整体设计思路从问题定义到方案选型1.1 先搞清楚我们要解决什么问题很多新手上来就写代码结果写到一半发现数据集不对、标签和图片对不上、模型输入尺寸错误最后调了两天 bug 才跑通。这类问题的根子在于一开始没把问题定义清楚。图像分类任务的标准描述是这样的给定一张图片模型从预设的类别集合中预测它属于哪一类。具体到本项目输入是运动鞋照片输出是品牌名比如 Nike、Adidas、New Balance。数据集准备好之后用 Pytorch 把图片读进来过一遍卷积神经网络拿到每个类别的概率分布取概率最高的那个作为最终的预测结果。这里有一个容易忽略的点模型输出的不是标签字符串而是每个类别的概率向量。这也是初学者最容易懵的地方——模型内部做的其实是一次数值计算把 3 通道的图片张量逐层卷积、池化、全连接最后输出一个长度为 N 的向量N 等于类别数向量里的每个值代表对应类别的置信度。Softmax 会把这些值归一化成概率但即使不显式调用 Softmax直接看 logits 的大小也能判断预测结果。1.2 为什么选 Pytorch 而不是其他框架项目标题里既然点名了 Pytorch我还是要说一下我在选型时的真实考量。Pytorch 最大的优势是动态图机制也就是说网络的每一层在运行时都是边算边构建的这意味着可以在循环里用 if 语句控制网络结构、在调试时用 print 查看任意中间变量形状这对排查问题无比友好。相比之下静态图框架需要先把图建好再执行调试体验会差一个量级。另一个很现实的原因是生态。Pytorch 官网上提供了 torchvision 这个视觉工具库里面预置了 ResNet、VGG、MobileNet 等主流模型结构和在 ImageNet 上预训练好的权重加载预训练模型只需要一行代码。我们做图像识别绝大多数情况下不需要从零训练一个网络站在预训练模型的肩膀上做迁移学习效果好、收敛快、对算力要求低这是目前视觉领域公认的高效方案。还有一个细节是社区资源的丰富程度。用 Pytorch 写图像分类遇到问题随手一搜就有大量讨论帖和源码可以参考官网文档也写得足够清楚API 设计直观从 tensor 操作到模型定义再到训练循环心智负担很小。对这些理由我的判断是Pytorch 是目前做图像分类项目最合适的选择没有之一。1.3 方案总览与核心流程整个项目按照这条线走的数据准备 → 数据预处理 → 模型搭建 → 训练调参 → 评估预测。数据准备阶段把拍摄好的运动鞋图片按品牌放到不同文件夹里数据预处理阶段用 torchvision 的 transforms 做缩放、裁剪、归一化模型搭建阶段加载 ResNet18 预训练权重并替换最后一层全连接训练阶段用交叉熵损失函数和 Adam 优化器迭代若干轮最后在验证集上统计准确率、绘制混淆矩阵拿几张没见过的图片做实际预测。这个方案不是唯一的比如模型可以换成更大的 ResNet50 或更轻量的 MobileNet优化器可以换成 SGD数据增强策略也可以更激进。但核心思想是一致的用预训练模型做特征提取器根据自己的任务微调最后几层既省时间又稳。下面我把每个环节的细节展开讲。2. 数据集准备与预处理决定模型上限的关键一步2.1 数据从哪来、目录怎么组织图像分类项目的模型训练代码其实非常标准化真正拉开效果差距的往往是数据。这次运动鞋识别项目的数据来源主要有两个渠道一部分来自公开的商品图片数据集另一部分是自己用手机拍的实拍图。我建议你也混合收集因为公开数据集里的图大多是白底、构图居中、光线均匀的“标准图”而实际使用场景中用户拍的图往往背景杂乱、角度刁钻、光线不佳只用标准图训练出来的模型在新场景下很容易翻车。数据集目录结构直接照着 torchvision 的 ImageFolder 约定来组织即可因为这个类会自动按照子文件夹名生成标签。目录是这样的data/ ├── train/ │ ├── nike/ │ │ ├── nike_001.jpg │ │ ├── nike_002.jpg │ │ └── ... │ ├── adidas/ │ │ └── ... │ └── new_balance/ │ └── ... └── val/ ├── nike/ ├── adidas/ └── new_balance/train 和 val 按 8:2 左右划分。划分的时候要注意一个新手很容易犯的错误不要在分类层面划分要按图片随机划分确保同一个鞋款的不同照片不会全部落在训练集或验证集否则验证结果会有虚高的成分。我习惯写一个小脚本先把所有图片路径打乱再按比例切分到两个目录里。图片数量这块我的建议是每个类别至少准备 100 张以上越多越好。如果只有几十张也不是完全不能跑手动加一些旋转、翻转、色彩抖动之类的数据增强也能救一救但泛化能力会明显受限。收集图像时尽量把不同背景、不同光线、不同角度的图都包含进来这样模型才有机会学到真正属于鞋子本身的特征。2.2 transforms 预处理为什么这么做读入图片后不能直接扔进网络因为预训练模型对输入有固定要求。ResNet 系列网络接受的是 224×224 的 RGB 图像同时要求输入的数值分布尽量接近 ImageNet 数据集的分布这样才能让预训练权重发挥最大作用。torchvision 的 transforms 提供了整套工具我用的组合是from torchvision import transforms train_transforms transforms.Compose([ transforms.RandomResizedCrop(224), 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_transforms 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]) ])这里每个 transform 都有它的道理。RandomResizedCrop 会随机裁剪图片的一部分再缩放到 224×224相当于模拟鞋子在画面中的不同大小和位置是一种非常有效的数据增强方式RandomHorizontalFlip 随机水平翻转让模型对鞋子的左右朝向不敏感ColorJitter 调整亮度、对比度和饱和度模拟不同光线条件下的拍摄效果。验证集不搞这些花活只做缩放和中心裁剪目的是用统一的预处理方式评估模型真实性能。Normalize 这一步很多人只抄不思考。为什么要用 0.485、0.456、0.406 这组均值因为这是 ImageNet 数据集的 RGB 三通道均值和方差预训练模型是在这些统计量下训练出来的输入数据对齐之后才能让模型“看到”和训练时一致的分布。如果自己从零训练模型用零均值一方差归一化也可以但既然用了预训练权重就遵循它原本的统计口径。2.3 DataLoader 的参数设置与数据增强的边界数据准备好之后用 DataLoader 包装有几个参数值得说。batch_size 控制每个批次送入模型的图片数量我这次设为 32如果你的显存比较紧张降到 16 或者 8 也没问题。shuffleTrue 是训练集必须开启的让每个 epoch 里图片顺序不同避免模型学到样本顺序的偶然而非真正的特征。num_workers 控制数据加载的子进程数Windows 下建议设为 0Linux 下可以设成 4 或 8这个参数设高了反而容易报错。数据增强不是越多越好这一点我吃过亏。有一段时间我把数据增强策略搞得非常激进加了随机旋转、随机擦除、高斯模糊结果训练集上 loss 下降变慢验证集准确率也没提升多少。后来我意识到数据增强的本质是正则化适度的扰动让模型见过更多变体、降低过拟合但扰动过头会让模型学不到稳定特征尤其是像运动鞋这种结构特征明确的对象过度旋转会改变鞋子在视觉上的合理性。常规的随机裁剪、翻转、色彩抖动已经足够应付大多场景先跑一版基线效果不满意再逐步增强策略。3. 模型搭建与训练配置站在巨人的肩膀上做迁移学习3.1 迁移学习的原理和为什么有效训练一个深层卷积网络需要海量数据和大量计算资源ImageNet 数据集有 128 万张图片、1000 个类别用多卡 GPU 也要训练好几天。我们手里的运动鞋数据最多不过几百张远不够从零训练一个深层网络。迁移学习解决了这个问题先在一个大规模数据集上练好一个模型让它学会提取通用视觉特征——比如边缘、纹理、颜色、形状结构——然后再把它搬到我们的任务上做微调。打个比方就容易理解了这相当于请了一位见过几百万张图片的“视觉实习生”他已经知道怎么识别物体的边缘、轮廓、材质现在只需要教他“这些特征组合在一起意味着这是 Nike 还是 Adidas”。他不需要从头学习怎么看图只需要在原有基础上快速适应新任务。计算机视觉领域几十年的经验表明底层卷积核学习到的是通用特征越靠近输出层的特征越跟具体任务相关所以迁移学习的常规做法是复用底层特征、重新训练顶层分类器。3.2 加载预训练模型并改造分类头torchvision 里加载预训练 ResNet18 就一行代码但需要注意的是最后一层全连接层的输出维度要改成我们的类别数。ResNet18 的最后一层全连接输入维度是 512默认输出 1000ImageNet 的类别数我们把输出改成 3 即可。下面是我用的完整代码import torch import torch.nn as nn from torchvision import models device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(pretrainedTrue) num_features model.fc.in_features model.fc nn.Linear(num_features, 3) model model.to(device)这里有两个地方要特别注意。第一是pretrainedTrue会自动下载 ImageNet 的预训练权重网络通畅的话会从官方地址下载如果下载速度很慢或者超时可以手动下载权重文件放到 torch 缓存目录路径一般是~/.cache/torch/hub/checkpoints/不同操作系统和用户目录会有差异报错信息里通常会写明期望的路径。第二是替换model.fc这个操作替换之后模型的输出维度就变成 3 了其他层的参数还是预训练好的。关于微调策略常见的有两种一是冻结特征提取层的参数只训练最后一层全连接二是所有层一起参与训练用较小的学习率微调。我这次选择的是全量微调因为运动鞋这类商品图片和 ImageNet 里的自然图像差异并不是特别大全量微调能更快适应新任务。如果你用的是非常小的数据集比如每个类别只有二三十张那冻结大部分层、只训练最后的分类头会更稳妥能降低过拟合风险。3.3 损失函数、优化器与学习率的选择逻辑损失函数用的是交叉熵损失这是多分类问题的标准选择。Pytorch 的nn.CrossEntropyLoss已经内置了 Softmax 计算所以模型的最后一层不需要额外接 Softmax预测时直接取 logits 最大值对应的索引作为类别即可。交叉熵损失衡量的是模型预测概率分布和真实标签分布之间的差异值越小说明预测越接近真实标签。从数值上看它等价于对真实类别的预测概率取负对数预测概率越接近 1损失越接近 0。优化器我选了 Adam初始学习率 0.001。Adam 自适应调整每个参数的学习率用起来省心不用太操心手调学习率的问题。如果你的训练过程出现 loss 震荡剧烈或者不下降的情况第一个该检查的就是学习率把学习率降到 0.0001 再试一轮往往症状就缓解了。SGD 配合 momentum 也是经典选择但需要更精细的学习率调度我在这种规模的入门项目上更推荐 Adam少一个调参变量跑通流程是首要目标。训练轮数我设了 30 轮配合StepLR每 10 轮把学习率乘以 0.1。这种学习率衰减策略的思路是训练前期用较大学习率迅速逼近最优区域后期用小学习率在最优解附近精细调整避免在最优解附近来回震荡。完整训练配置如下import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1)4. 完整训练流程与结果评估4.1 训练循环的写法与关键细节训练循环是整套代码的核心我把它拆成函数来写结构清晰也好复用。每个 epoch 里做的事可以归纳成四步前向传播计算输出和损失梯度清零反向传播计算梯度优化器更新参数。这里最容易漏掉的就是optimizer.zero_grad()如果忘了清零梯度会在多个 batch 上累积导致参数更新方向错误。Pytorch 的梯度默认是累积的这是它的设计特性但训练时我们通常希望每个 batch 独立更新一次。训练阶段要用model.train()切换到训练模式这个模式下 BatchNorm 层会使用当前 batch 的统计量Dropout 层会生效。验证阶段要用model.eval()切换到评估模式BatchNorm 使用训练阶段累计的统计量Dropout 被关闭。如果忘了切换BatchNorm 的行为会异常验证准确率甚至可能掉到随机水平。这个坑我见过太多人踩了代码里务必养成成对写的习惯。def train_one_epoch(model, dataloader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in dataloader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_accloss.item()的作用是从 GPU 上的 tensor 取出 Python 数值方便记录和绘图。torch.max(outputs, 1)返回每行最大值及其索引索引就是预测的类别编号。注意loss.item()乘上images.size(0)是为了计算整个 epoch 的总损失因为 loss 本身是一个 batch 的平均值要还原成总量才能正确求全局平均。同样的逻辑用在计算准确率上用correct / total而不是每个 batch 单独算准确率再平均两种方法在数据量整除时结果一致但前者的数学含义更严谨。4.2 验证评估与模型保存每个 epoch 训练完之后跑一遍验证集记录验证集上的准确率和损失用来判断模型是在“进步”还是“过拟合”。验证阶段不需要计算梯度用torch.no_grad()包住可以省下大量显存和计算时间这是一个几乎不影响代码可读性但显著提升效率的好习惯。完整代码如下def evaluate(model, dataloader, criterion, device): model.eval() val_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in dataloader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) val_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss val_loss / total accuracy correct / total return avg_loss, accuracy我的习惯是每个 epoch 结束时把验证集准确率最高的模型权重保存下来而不是简单保存最后一轮的权重。因为模型可能在中间的某个 epoch 泛化能力最好最后一轮不一定是最优的。保存整个模型还是只保存权重取决于项目用途。只保存权重更省空间也更灵活best_acc 0.0 for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc evaluate(...) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fEpoch {epoch1}/{num_epochs} | Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}) # 预测时加载权重 model.load_state_dict(torch.load(best_model.pth))打印日志的格式我故意调整了很多次最后固定成上面对齐方式。每行输出训练集和验证集的 loss 与准确率横看一个 epoch 内的表现竖看整个训练的走势。跑完 30 轮以后把日志拉出来如果训练集准确率一直涨但验证集准确率开始回落过拟合的特征就非常明显了。4.3 参数调整的实际经验跑完第一轮训练后我看到的验证集准确率大概在 85% 左右。这个成绩对于迁移学习来说算正常水平说明预训练模型已经学到了足够好的视觉特征。接下来花在调参上的时间主要在几个方向一是把学习率从 0.001 调整为前几轮先用 0.01 加速收敛、后面衰减到 0.0001二是增加了数据增强的强度比如把 ColorJitter 的范围调大了一些三是把训练轮数从 30 增加到 50配合早停策略当验证集准确率连续 5 轮没有提升就提前结束训练。最终模型的验证集准确率达到了 92% 左右。这三个措施里收益最明显的是早停和学习率调度数据增强的调整反而收益不大。这也印证了一个经验模型效果提升的优先级应该是先确保数据质量再调训练策略最后才动模型结构。大多数情况下与其纠结从 ResNet18 换成 ResNet50不如先把已有的方案打磨到极致。5. 常见问题与排查技巧实录5.1 过拟合的判断与应对项目里出现过几次过拟合的苗头表现是训练集准确率在 97% 以上但验证集准确率停滞在 88% 左右不上去了并且训练 loss 稳步下降的同时验证 loss 反而在上升。这种情况说明模型开始“死记硬背”训练集里的具体样本而不是抽象出通用的“运动鞋特征”。应对手段从轻到重排列的话第一个是增强数据增强的强度让模型在训练中看到更多变化第二个是加 Dropout 层但 ResNet 架构里本身没放 Dropout要加只能在全连接层后面加改动相对麻烦第三个是冻结更多底层网络层只微调后面几层限制模型的表达能力第四个是收集更多数据。我在这个项目里主要用了第一个手段加大 ColorJitter 的范围之后验证集准确率从 88% 提到了 91% 左右效果立竿见影。5.2 Pytorch 环境与运行时的典型报错处理先说环境安装的问题。Pytorch 官网提供了一键安装命令Conda 和 Pip 两种方式都支持但最常出问题的是 CUDA 版本不匹配。我的建议是先确认显卡驱动支持的最高 CUDA 版本再到官网选择对应版本的安装命令。如果只是做 CPU 推理直接装 CPU 版就行省掉一堆驱动和运行时排查的麻烦。用torch.cuda.is_available()可以快速验证 CUDA 是否可用。运行时最常见的报错是显存不足CUDA out of memory这时优先把 batch_size 减半其次减小图片分辨率最后再考虑换更轻量的模型。还有一个容易踩的坑是 DataLoader 的num_workers在 Windows 上不设 0 会报错报错信息提到BrokenPipeError十有八九是这个原因。另外Pytorch 现在从 1.x 升级到 2.x 之后部分旧 API 有兼容性调整如果跑别人的旧代码报错优先去看是否有torchvision.transforms导入路径变化或者pretrained参数被弃用的提示。5.3 图片质量与推理效果的注意点模型训练完以后我拿了手机实拍的照片做测试结果发现某些角度的鞋子识别准确率明显下降。后来分析了一下主要是拍摄角度问题——有些图片是俯拍的有些是侧面拍的爆炸视角下鞋子的轮廓特征和训练数据中的标准图差异很大。解决办法是给训练数据增加更多不同角度的图同时可以用 RandomResizedCrop 配合较大的裁剪范围来模拟不同的构图方式。还有一个细节很多人不重视不要忘了在预测时对输入图片应用和验证集一样的预处理。如果训练时做了 Resize、CenterCrop 和 Normalize预测时跳过了这些步骤模型拿到的就是“分布外”的数据预测结果自然不可信。我曾见过有人训练时用 224×224 输入预测时直接传入原始尺寸的图片导致模型输出完全不合理的概率分布检查了很久才发现是这个低级错误。做预测的代码一定要复用验证集的 transforms 定义。排查问题的过程中我把遇到过的典型问题整理成了一个速查表供参考现象可能原因排查与解决方法训练 loss 不下降学习率过大或过小调到 0.0001 到 0.01 区间重试验证集准确率远低于训练集过拟合增强数据增强、冻结更多层、加 DropoutCUDA out of memorybatch_size 或输入尺寸过大减半 batch_size降低图片分辨率DataLoader 报 BrokenPipeErrorWindows 上 num_workers 设置问题把 num_workers 设为 0训练和验证准确率都很低数据标签错误或预处理问题检查标签对应关系确认 Normalize 参数预测结果完全不对预测时未做预处理复用验证集的 transforms 流程这个项目跑下来的整体感受是用 Pytorch 做图像分类真正难的不是框架的 API而是对数据质量和流程细节的把控。预训练模型把特征提取的门槛降得很低一个在 ImageNet 上练过的 ResNet18 拿到几万张运动鞋图片上微调效果已经很可观了。如果你想要进一步提升准确率可以从几个方向入手把 ResNet18 换成 EfficientNet 或 ConvNeXt增加训练数据量或者针对运动鞋的局部特征做更细粒度的分类。这些我在后续的系列里会继续分享先把这个基础流程吃透后面搭什么模型都不慌。