很多刚接触 PyTorch 的朋友都会问第一个练手项目选什么我的答案一直是同一个——猫狗二分类。它看起来只是简单地把图片分成猫和狗两类但背后涉及数据加载、图像预处理、模型构建、训练调参、指标评估和模型保存是一条非常完整的深度学习流程。我最初完整跑通的项目就是它后来再看很多相似分类任务代码骨架几乎没有变过。所以这篇内容不是想讲花哨技巧而是把这条路上最容易踩坑、最值得想清楚的地方摊开来说。如果你正准备学 PyTorch或者已经会一点基础但没完整做过分类项目这篇内容应该能帮你省不少时间。1. 猫狗二分类项目拆解新手为什么绕不开它1.1 为什么几乎所有教程都选猫狗分类练手猫狗分类最吸引人的一点是“数据不难”。自己用手机拍一些家里宠物照片或者找一些公开的小样本猫狗图片就能凑出一个训练集。它不像医疗影像、卫星图那样需要专业标注背景猫和狗这两个类别本身有相对明显的纹理、轮廓和姿态差异模型只要学到一些有效特征就很容易看到效果。这种“快速获得正反馈”的体验对入门者非常重要。但这个任务也没有简单到“闭眼就能做好”。猫狗图片里有大量背景干扰有遮挡、暗光、视角倾斜甚至有些猫长得像狗、有些狗看起来像猫。真正把验证准确率做到 95% 以上需要认真处理数据划分、数据增强、模型容量和过拟合控制。它刚好卡在“不难但不能乱来”的位置最适合用来建立完整的工程意识。另一个原因是这个项目能自然延伸。猫狗二分类是一套标准图像分类流水线换成猫狗三分类、场景分类、菜品分类改动无非是数据目录和最后一层输出节点数。把这次项目的代码保存好以后遇到新的分类需求基本能在半小时内跑通 baseline。我身边的很多朋友都是先做完猫狗二分类之后才敢去碰更复杂的检测和分割任务。1.2 技术选型拆解PyTorch、CNN 和迁移学习选择 PyTorch 的理由在我看来主要是三点动态图调试方便、生态完整、社区资料多。动态图的意思是你的模型每跑一步中间结果都是实实在在能打印、能断点检查的张量。新手最怕的就是“模型跑通了但不知道里面在算什么”PyTorch 允许你在前向传播里任意插入打印语句这对理解张量形状变化非常有帮助。图像分类不需要一开始就上 TransformerCNN 依然是性价比最高的起点。卷积层相当于在图像上滑动一组可复用的模板局部连接和参数共享让它比全连接层更适合捕捉像素之间的空间关系。你可以把卷积核想象成一套“特征扫描器”底层扫描边缘、颜色块高层组合出眼睛、耳朵、尾巴这类语义部件。对猫狗分类这种中小规模任务一个轻量 CNN 完全有能力学到不错的特征。迁移学习则是锦上添花。PyTorch 官方 torchvision 里提供了很多在大型图像数据集上预训练好的模型比如 ResNet18。预训练模型已经学会了大量通用视觉特征我们只需要把最后一层全连接替换成自己的二分类输出头再微调几轮通常比从零训练快得多也稳得多。数据量越少迁移学习的优势越明显。后面我会把“从零搭 CNN”和“迁移学习微调”两条路线都写出来你可以在同一份数据上做对比感受非常直观。2. 准备阶段最容易忽略的细节环境、数据划分与增强2.1 环境准备与硬件选型我建议的 Python 环境是 3.10 或 3.11搭配 PyTorch 2.x 和配套的 torchvision。安装 PyTorch 时最重要的一点是确认版本匹配尤其是 GPU 版本需要与你的驱动和 CUDA 版本对应。很多人一上来就报错“CUDA not available”大概率不是显卡不行而是安装了 CPU 版或驱动版本太老。在终端里输入python -c import torch; print(torch.cuda.is_available())如果能输出True环境才算合格。硬件方面CPU 也能跑猫狗分类但训练会非常痛苦。一个几十万参数的小 CNN 在 CPU 上跑 20 轮可能需要几个小时而在独立显卡上可能只要十几分钟。如果你暂时没有 GPU我建议先把数据量缩小、batch size 调小、图片分辨率降到 96×96先验证整套流程能跑通。等到换到 GPU 上再跑全量数据会轻松很多。torchvision 一定要装因为除了模型定义它还提供transforms图像处理工具和datasets.ImageFolder数据加载接口。很多新手以为只需要 torch 就够结果处理图片时到处手写循环既慢又容易出错。数据加载这块用 torchvision 自带组件是成熟且省力的做法。2.2 数据集整理从零开始把图片喂给模型我最推荐的目录结构是下面这种它可以直接配合datasets.ImageFolder使用data/ train/ cat/ dog/ val/ cat/ dog/ImageFolder会自动把train下的一级子目录当作类别图片按文件名逐个读取并生成从 0 开始递增的类别索引。注意这里有一个隐藏规则类别顺序是按子目录名的字母序自动排序的。也就是说cat会被分到索引 0dog会被分到索引 1。如果你把子目录命名为dog和cat则顺序反过来。后面做预测时一定要先确认这个映射关系否则模型输出 0 和 1 时你会分不清到底哪个是猫。如果手里的图片没有现成的验证集记得自己做随机划分。我习惯先把所有图片路径读取到一个列表里按类别分层采样例如 80% 放训练、20% 放验证。分层的意思是保证猫和狗在两部分里比例相近避免出现训练集全是猫、验证集全是狗这种极端情况。划分完成后可以用shutil把文件复制到对应目录。这里不要偷懒直接把全部图片丢进去训练没有验证集你根本无法判断模型是否过拟合。2.3 数据增强不要让模型死记硬背数据增强是猫狗分类里提升效果最直接的手段。我的常用训练增强包括随机水平翻转、随机旋转 10 度、随机亮度饱和度抖动、随机缩放裁剪。验证阶段不能加这些随机变换只需要统一尺寸、转张量、归一化。from torchvision import transforms train_transforms transforms.Compose([ transforms.Resize((128, 128)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(10), 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((128, 128)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])归一化的均值方差为什么要用0.485, 0.456, 0.406这一组因为它们是在大型公开图像数据集上统计得到的 RGB 三通道均值方差。无论你用 PyTorch 官方预训练模型还是自己训练统一使用这组参数都能让训练更稳定。如果你用自己完全从零训练的模型理论上可以重新统计自己数据的均值方差但大多数情况下沿用这组经验值就够了。数据增强的本质是“给模型多找点事做”。如果每张猫图都是同一个角度、同一亮度模型很可能把背景误当成猫的特征而不是真正学会猫本身的轮廓。做了增强之后模型被迫学习不变性不管猫向左还是向右不管光线偏亮还是偏暗都要认出这是猫。但增强也不是越猛越好。我见过有人把旋转角度设成 90 度结果训练集里出现大量倒立猫狗模型反而学乱了。像猫狗这种非旋转不变的任务控制在 10 度到 15 度比较稳妥。3. 模型搭建与核心训练代码从零实现到迁移学习3.1 从零写一个轻量 CNN 全流程虽然可以直接用迁移学习但我强烈建议新手先自己写一个小 CNN认真跑一遍训练循环再上预训练模型。自己写模型时最核心的是想清楚张量形状。下面是一个适用于 128×128 输入的轻量网络import torch import torch.nn as nn import torch.nn.functional as F class CatDogNet(nn.Module): def __init__(self, num_classes2): super().__init__() self.conv1 nn.Conv2d(3, 32, kernel_size3, padding1) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.conv3 nn.Conv2d(64, 128, kernel_size3, padding1) self.pool nn.MaxPool2d(2, 2) # 输入 128x128经过三次池化后变成 16x16 self.fc1 nn.Linear(128 * 16 * 16, 256) self.fc2 nn.Linear(256, num_classes) def forward(self, x): x self.pool(F.relu(self.conv1(x))) # 128 - 64 x self.pool(F.relu(self.conv2(x))) # 64 - 32 x self.pool(F.relu(self.conv3(x))) # 32 - 16 x x.view(x.size(0), -1) x F.relu(self.fc1(x)) x self.fc2(x) return x这个网络有三个卷积层每层后面接 ReLU 和 2×2 最大池化。为什么用 padding1因为这样卷积不会改变特征图尺寸尺寸收缩完全由池化负责展平前的形状好计算。你也可以用 padding0 并在代码里手动算尺寸但新手很容易算错保持 padding1 能让流程更顺。fc1输入维度是 128×16×16对应最后一个池化输出展平后的长度这个数字要特意算清楚。如果你的输入尺寸不是 128这段代码里的 16 就要相应改变。3.2 损失函数、优化器与学习率设置分类任务最常用的损失函数是nn.CrossEntropyLoss()。它内部把 softmax 和交叉熵计算整合在了一起所以网络最后一层不需要额外手动加 softmax。猫狗二分类的数据标签通常是 0 或 1模型输出是一个长度为 2 的向量比如[2.1, -0.5]交叉熵损失会自动把它转换成概率并计算 loss。优化器我推荐先使用 Adam学习率设 0.001。Adam 几乎不需要太多手动调整就能跑得不错对新手非常友好。它有自适应学习率相当于帮你自动缩放每一步更新的幅度。等你已经能手调 SGD 的动量、学习率衰减再换也不迟。训练主循环的骨架如下model CatDogNet() criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr0.001) for epoch in range(20): model.train() total_loss 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss / len(train_loader):.4f})注意到每个 batch 我都没有写“梯度清零会怎样”的废话但这里确实是新手踩坑点。PyTorch 的梯度是累积的不清零的话每个 batch 的梯度会叠加在一起相当于把好几张图的信息混在一起更新参数。所以optimizer.zero_grad()必须放在loss.backward()之前。常规顺序就是前向算 loss、清零梯度、反向传播、更新参数。3.3 迁移学习站在预训练模型肩膀上如果你想快速看到更高准确率可以迁移学习。以 ResNet18 为例torchvision 里提供了在大型图像数据集上预训练好的版本。我们只需要把原来的最后一层全连接替换掉import torchvision.models as models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) num_features model.fc.in_features model.fc nn.Linear(num_features, 2)注意两点。第一ResNet18 原本的输入尺寸是 224×224所以你如果用迁移学习数据增强里的Resize最好改成(224, 224)而不是 128。第二如果只替换全连接层前面的卷积层仍然保留预训练参数。训练初期可以先冻结这些特征层只让全连接层更新for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True这样做的原因是预训练特征已经足够通用我们只需要让分类头学习“如何组合这些特征”。冻结参数能大幅减少计算量也降低过拟合风险。如果数据量比较大或者冻结之后验证准确率不够高再解冻部分特征层做全模型微调。解冻后学习率建议调小一点比如 0.0001避免预训练权重被破坏。4. 训练过程实录超参数、曲线和过拟合怎么处理4.1 超参数设置与实验记录表格训练时我习惯把关键超参数记在实验表格里。不要相信自己的记忆力调参过程中经常会有“昨天我改了什么来着”的尴尬。下面是我在模拟项目 X 中最常用的初始超参数超参数从零训练小 CNNResNet18 迁移学习输入尺寸128×128224×224Batch Size3232优化器AdamAdam初始学习率0.0010.0001训练轮数2010数据增强水平翻转、旋转、颜色抖动水平翻转、旋转输出层节点22同一份数据下从零训练的小 CNN 大约需要 20 轮才能稳定验证准确率一般在 88% 到 93% 左右。ResNet18 迁移学习可能在 5 轮内就追上来10 轮之后很容易超过 95%。这个差距非常直观地体现了预训练权重的作用。如果你是首次跑这个项目建议把两种方案都跑一遍形成自己的对比数据会远比只看别人的结论有用。4.2 训练曲线怎么看过拟合怎么治训练曲线是诊断模型状态的第一手来源。我见过最典型的曲线是训练 loss 持续下降但验证准确率在第 6 轮之后不再上升甚至开始下跌。这说明模型开始“死记硬背”训练集了也就是过拟合。你可以在代码里记录每轮的训练 loss 和验证准确率训练结束后用 matplotlib 画出来两条线的距离一旦拉得很大就该警惕了。处理过拟合的手段按优先级排列增加数据增强、使用迁移学习、减小模型容量、加 Dropout、早停。早停是最容易实现的手段核心是保存验证准确率最高的那一轮权重而不是用最后一轮的权重best_acc 0.0 for epoch in range(num_epochs): # 训练和验证 val_acc evaluate(model, val_loader) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth)如果验证准确率连续多轮不涨就可以提前结束训练。很多人会把全部 20 轮跑完但最优模型可能在第 12 轮就已经出现后面几轮纯粹在浪费时间甚至伤害泛化能力。早停不是偷懒而是避免无效计算的实用策略。4.3 迁移学习与从零训练的对比结果从零训练的好处是你对模型内部结构的理解更扎实。你能清楚看到每个卷积层输出的特征图逐渐从“边缘纹理”变成“眼睛耳朵”这对建立深度学习直觉非常重要。缺点是对数据量要求更高训练时间更长稍不注意就容易过拟合。迁移学习则是“站在前人的肩膀上”。ResNet18 的预训练特征已经具备很强的图像理解能力我们只需要微调最后几层。它更适合真实场景中数据量有限、想要快速上线 baseline 的情况。代价是模型文件更大加载稍慢如果你部署环境非常受限可能需要蒸馏或量化。我个人的建议是学习阶段从零训练一个模型项目阶段直接迁移学习。两者并不矛盾。等你两类方案都跑过一遍再回头看“为什么要用预训练模型”这类问题会有非常具体的答案。5. 评估与预测模型做出来到底能不能用5.1 用 accuracy 和混淆矩阵衡量真实水平训练过程里我一般用验证集准确率作为主要指标因为它简单直观。猫狗二分类样本均衡准确率没有太大欺骗性。但如果以后遇到类别不平衡的数据只看准确率会误判。比如 99% 是猫、1% 是狗模型全预测猫也能有 99% 准确率但狗这个类别完全没有被识别。所以学会看混淆矩阵很重要。混淆矩阵把预测结果分成四类真正例、假正例、真负例、假负例。对猫狗分类来说可以看“狗被错分成猫”和“猫被错分成狗”具体有多少张。生成方法很简单先遍历验证集收集所有预测标签和真实标签然后用sklearn.metrics.confusion_matrix输出矩阵。如果发现狗被错分成猫的比例特别高说明模型在狗类特征上学得不够需要补充狗的样本或增强狗的图片多样性。5.2 单张图片预测的完整流程模型训练完最终要能用起来。推理流程和训练不同不需要计算梯度也不需要优化器。我只保留一个干净的单图预测函数def predict_image(image_path, model, device, class_names): image Image.open(image_path).convert(RGB) image val_transforms(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): outputs model(image) probs torch.softmax(outputs, dim1) pred_idx torch.argmax(probs, dim1).item() return class_names[pred_idx], probs[0][pred_idx].item()这个函数有四个关键点。第一convert(RGB)是为了防止灰度图或带透明通道的 PNG 导致输入通道数不对。第二unsqueeze(0)是在 batch 维度上扩一维因为模型期望的输入形状是(batch, channel, height, width)单张图也需要加一个维度。第三推理时用model.eval()和torch.no_grad()前者关闭 Dropout 和 BatchNorm 的训练状态后者停止构建计算图省内存也提速。第四容易忽略的是class_names顺序它应该和训练时ImageFolder的类别索引完全对应比如[cat, dog]。5.3 模型保存与加载的常见坑保存模型最稳妥的方式是保存state_dict也就是模型参数字典而不是用torch.save(model, ...)保存整个模型。前者体积更小、更稳定加载时只需先创建相同结构的模型再载入参数。加载代码看起来简单却有几个容易翻车的细节model CatDogNet() state_dict torch.load(best_model.pth, map_locationcpu) model.load_state_dict(state_dict) model.to(device)map_locationcpu的作用是如果模型是在 GPU 上训练的而你现在只想要 CPU 加载它能帮你做张量设备映射避免“CUDA 不可用”之类的报错。另外加载后别忘了调用model.eval()。如果直接拿去预测模型可能还处于训练模式BatchNorm 和 Dropout 的行为不一样结果会变差。我们前面所有训练、验证、保存、加载的代码都是在和model.train()、model.eval()这两个状态打交道建议真正理解它们的区别再去部署。6. 新手实战排雷我踩过的坑和解决清单6.1 数据加载与图像尺寸相关的坑第一个坑是类别顺序。刚开始跑通时我做过一次预测发现输出结果显示“狗”但图片里明明是猫。排查到最后才发现ImageFolder的类别索引不是按我直觉里的“猫在 0、狗在 1”而是需要显式查看dataset.class_to_idx。这一步在训练前打印出来能省掉后面一大段困惑。第二个坑是输入尺寸不一致。如果训练集里大部分图片是 128×128但某张图片是 256×256而你使用的 transform 漏掉了 Resize就会在读 batch 时报 shape 不一致。transforms.Compose里把Resize((128, 128))放在最前面基本能避免这类问题。千万不要用自己的裸循环去读一张张图再 resize正确做法是让 DataLoader 配合 transform 在加载阶段直接统一。第三个坑是归一化参数不一致。这句话我重复了很多次因为见过太多人在迁移学习时把 Normalize 设置成mean0,std1等于没归一化训练初期 loss 迟迟不降。如果使用预训练权重就一定要沿用对应的归一化参数。6.2 显存不足与随机种子问题跑迁移学习时224×224 输入加上较大 batch size很容易碰到“CUDA Out of Memory”。常见对策是降低 batch size、减小输入尺寸、用完的中间变量及时释放、验证时不用torch.no_grad()的话会白白占用大量显存。还有一个不太起眼但很实用的建议训练循环里不要反复调用images.to(device)之后再创建新的冗余张量尽量复用变量名。随机种子问题则是复现实验的关键。很多人跑两次结果不一样会误以为模型不稳定实际上是没有固定随机种子。固定方式可以在训练脚本开头加上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.3 避坑策略与个人实操心得如果让我给新手一个最值得坚持的策略那就是“先用小数据把流程跑通再上全量数据”。我见过很多朋友一上来就加载全部图片然后折腾半天发现代码报错报错信息被淹没在大量数据读入里根本看不清。正确的做法是先用每个类别几十张图组成一个小训练集跑一个 epoch确认前向反向、指标打印、模型保存都没问题再换全量数据训练。这一步能筛掉大部分低级错误。另一个心得是模型没效果先怀疑数据流不要上来就调结构。把一张图片从硬盘读到 DataLoader再前向输出结果中间每一步都可以打印出来检查。看看图片像素值范围对不对、标签对不对、网络输出形状对不对。很多问题只要把张量形状打出来一眼就能定位。我在跑猫狗分类时也走过“改网络、改损失、改学习率”的弯路最后发现只是 transform 里漏了一个 Resize导致通道数和预期不一致。从那以后我再遇到训练不收敛第一件事永远是查数据加载链路。这套从数据准备、模型搭建、训练调参到评估预测的流程今天看会觉得有些基础但它的确是我认为最扎实的 PyTorch 分类入门路径。你把它完整跑下来之后再去看更复杂的图像任务会发现很多套路都是相通的。最后分享一个实用习惯每次实验保存时把数据增强方式、超参数、验证准确率都写进同名的文本文件。这样几周后再回来翻实验记录也能清楚知道自己当时做了什么决策而不是对着文件夹里一堆model_final_v2.pth发呆。
