很多刚接触 PyTorch 的朋友都会问:第一个练手项目选什么?我的答案一直是同一个——猫狗二分类。它看起来只是简单地把图片分成猫和狗两类,但背后涉及数据加载、图像预处理、模型构建、训练调参、指标评估和模型保存,是一条非常完整的深度学习流程。我最初完整跑通的项目就是它,后来再看很多相似分类任务,代码骨架几乎没有变过。所以这篇内容不是想讲花哨技巧,而是把这条路上最容易踩坑、最值得想清楚的地方摊开来说。如果你正准备学 PyTorch,或者已经会一点基础但没完整做过分类项目,这篇内容应该能帮你省不少时间。
1. 猫狗二分类项目拆解:新手为什么绕不开它
1.1 为什么几乎所有教程都选猫狗分类练手
猫狗分类最吸引人的一点是“数据不难”。自己用手机拍一些家里宠物照片,或者找一些公开的小样本猫狗图片,就能凑出一个训练集。它不像医疗影像、卫星图那样需要专业标注背景,猫和狗这两个类别本身有相对明显的纹理、轮廓和姿态差异,模型只要学到一些有效特征,就很容易看到效果。这种“快速获得正反馈”的体验,对入门者非常重要。
但这个任务也没有简单到“闭眼就能做好”。猫狗图片里有大量背景干扰,有遮挡、暗光、视角倾斜,甚至有些猫长得像狗、有些狗看起来像猫。真正把验证准确率做到 95% 以上,需要认真处理数据划分、数据增强、模型容量和过拟合控制。它刚好卡在“不难但不能乱来”的位置,最适合用来建立完整的工程意识。
另一个原因是这个项目能自然延伸。猫狗二分类是一套标准图像分类流水线,换成猫狗三分类、场景分类、菜品分类,改动无非是数据目录和最后一层输出节点数。把这次项目的代码保存好,以后遇到新的分类需求,基本能在半小时内跑通 baseline。我身边的很多朋友,都是先做完猫狗二分类之后,才敢去碰更复杂的检测和分割任务。
1.2 技术选型拆解:PyTorch、CNN 和迁移学习
选择 PyTorch 的理由,在我看来主要是三点:动态图调试方便、生态完整、社区资料多。动态图的意思是你的模型每跑一步,中间结果都是实实在在能打印、能断点检查的张量。新手最怕的就是“模型跑通了但不知道里面在算什么”,PyTorch 允许你在前向传播里任意插入打印语句,这对理解张量形状变化非常有帮助。
图像分类不需要一开始就上 Transformer,CNN 依然是性价比最高的起点。卷积层相当于在图像上滑动一组可复用的模板,局部连接和参数共享让它比全连接层更适合捕捉像素之间的空间关系。你可以把卷积核想象成一套“特征扫描器”:底层扫描边缘、颜色块,高层组合出眼睛、耳朵、尾巴这类语义部件。对猫狗分类这种中小规模任务,一个轻量 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会被分到索引 0,dog会被分到索引 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(brightness=0.2, contrast=0.2, saturation=0.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_classes=2): super().__init__() self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) 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 最大池化。为什么用 padding=1?因为这样卷积不会改变特征图尺寸,尺寸收缩完全由池化负责,展平前的形状好计算。你也可以用 padding=0 并在代码里手动算尺寸,但新手很容易算错,保持 padding=1 能让流程更顺。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(), lr=0.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(f"Epoch {epoch+1}, 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(weights=models.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 中最常用的初始超参数:
| 超参数 | 从零训练小 CNN | ResNet18 迁移学习 |
|---|---|---|
| 输入尺寸 | 128×128 | 224×224 |
| Batch Size | 32 | 32 |
| 优化器 | Adam | Adam |
| 初始学习率 | 0.001 | 0.0001 |
| 训练轮数 | 20 | 10 |
| 数据增强 | 水平翻转、旋转、颜色抖动 | 水平翻转、旋转 |
| 输出层节点 | 2 | 2 |
同一份数据下,从零训练的小 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, dim=1) pred_idx = torch.argmax(probs, dim=1).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_location="cpu") model.load_state_dict(state_dict) model.to(device)map_location="cpu"的作用是,如果模型是在 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 设置成mean=0,std=1,等于没归一化,训练初期 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(seed=42): 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发呆。