PyTorch猫狗图像分类实战：从环境配置到模型训练完整指南

简介：这是一份面向深度学习初学者与有一定基础的开发者的PyTorch猫狗图像分类实战教程，提供从项目背景、数据增强、轻量级CNN搭建到训练评估与部署的完整流程。内容以中文讲解配合可直接复制运行的Python代码，覆盖随机裁剪、水平翻转、归一化等数据增强技巧，以及卷积层、Dropout、Adam优化等模型训练要点，适合个人自学、课堂教学或作为技术博客的写作模板。资源共1个docx文档，压缩包约19KB，内容结构紧凑，便于快速查阅与代码复用。目前已有152人学习，是入门计算机视觉分类任务的高性价比参考资料。

1. 用 PyTorch 做猫狗图像分类：为什么这是深度学习入门的最佳实战项目

图像分类是深度学习里最成熟、也最能讲清完整链路的方向，而猫狗分类正好是这类项目里数据最干净、模型收敛最快、结果最容易验证的一个。用 PyTorch 搭一个猫狗分类器，你会在同一个项目里走完数据准备、模型定义、训练、调参、评估、推理的全部流程，这套流程换到工业质检、遥感识别、医疗影像上，骨架是通用的。适合正在学深度学习但还没完整跑通一个训练闭环的初学者，也适合想快速验证 PyTorch 环境是否可用的从业者。

我见过太多人一上来就啃复杂目标检测或者 Transformer 模型，结果光环境就折腾两三天，最后连 loss 下降都没看到就放弃了。猫狗分类的工程复杂度刚好卡在一个舒服的位置——模型用现成分类网络，数据集是公开的小图集，本地 CPU 也能勉强跑、有 GPU 就更顺手。这篇笔记直接把我的完整做法和踩过的坑写出来，从环境搭建到训练完成，照着抄就能跑通。

2. 先立住 PyTorch 环境：装错版本会让整个项目卡死在第一步

2.1 环境组合的核心：CUDA、PyTorch、Python 三者必须匹配

做猫狗分类这类中小型数据集任务，环境问题的排查成本往往比模型调参还高。最常见的情况是：conda 里建好了环境，pip 也装了 torch，但一跑torch.cuda.is_available()就返回 False，GPU 完全没被识别到。这个问题的根源几乎都是 PyTorch 版本和 CUDA 驱动版本不匹配。

我的建议是不要自己凭感觉去 pip 安装。PyTorch 官方给出了很明确的安装矩阵，先搞清楚自己的显卡驱动支持到什么 CUDA 版本，再去选 PyTorch 的对应轮子。NVIDIA 驱动的 CUDA 版本在命令行里执行nvidia-smi就能看到，右上角就是当前驱动支持的最高 CUDA 版本。然后去 PyTorch 官网找对应组合，比如驱动支持 CUDA 11.8，就装 cu118 那个版本对应的 torch 安装命令。

# 先创建一个干净的 conda 环境，Python 版本选 3.9 或 3.10 都行 conda create -n catsdogs python=3.9 -y conda activate catsdogs # 安装 PyTorch，这里以 CUDA 11.8 版本为例 # 注意：cu118 表示这个轮子是用 CUDA 11.8 编译的 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

这段命令里最关键的是最后的--index-url参数。默认的 pip 源装的是 CPU 版，所以很多人装完以后以为自己有 GPU 加速，实际上跑的是 CPU。装完之后立刻验证一下：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

如果输出True并且能看到显卡型号，环境就准备好了。如果输出False，去看看是不是 conda 环境里本来就有一个 CPU 版的 torch 被优先识别了，pip list | grep torch看一下版本号有没有带+cu后缀。

2.2 Linux 服务器和 Windows 本机的环境差异

猫狗分类这种项目，数据量不大，放在 Linux 服务器上跑和放在 Windows 本机上跑差别不小。Linux 服务器一般都有 NVIDIA 驱动和 CUDA 工具链的基础配置，装 PyTorch 的步骤就是上面那两行，翻车的概率低。

Windows 本机的坑更多集中在 WSL 环境上。很多人用 WSL 里的 Linux 子系统配 PyTorch，但 WSL 默认不共享 Windows 的 GPU 驱动，必须先在 Windows 端装好 NVIDIA 的 WSL 驱动，再在 WSL 里执行nvidia-smi确认能看到 GPU。还有一个容易忽略的细节：WSL 里安装 CUDA 相关组件时要选 WSL-Ubuntu 版本而不是 Linux 通用版本。

纯 CPU 环境也不是不能跑。猫狗分类数据集一张图几百 KB，用 ResNet 这种模型训练一个 epoch 在 CPU 上可能要十几分钟，但如果只是验证代码逻辑、跑通训练流程，用一个小数据集或者只训练几个 epoch，CPU 完全够用。我自己习惯的做法是：先在一小部分数据上把完整流程跑通，再切到 GPU 上跑全量，这样排错效率最高。

3. 数据准备是整个项目的隐形重头戏：目录结构、预处理与数据划分

3.1 猫狗数据集的目录安排：用 ImageFolder 直接读是最省力的方案

PyTorch 的torchvision.datasets.ImageFolder是一个非常好用的数据加载工具，它要求数据目录按下层文件夹结构组织，每个子文件夹名就是类别名，文件夹里的图片就是该类别的样本。猫狗分类的目录结构可以这样搭：

data/ |-- train/ | |-- cat/ | | |-- cat.0.jpg | | |-- cat.1.jpg | |-- dog/ | | |-- dog.0.jpg | | |-- dog.1.jpg |-- val/ | |-- cat/ | | |-- cat.100.jpg | |-- dog/ | | |-- dog.100.jpg

这样组织好后，两行代码就能把数据和标签同时读出来。我一般不自己写 Dataset 类来读猫狗数据，因为 ImageFolder 已经把文件名映射标签这个活做完了，自己写反而容易在标签顺序上翻车。

from torchvision import datasets, transforms # 训练集和验证集用不同的预处理方式 train_dataset = datasets.ImageFolder( root='data/train', transform=transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), # 随机翻转，数据增强，提升泛化 transforms.ToTensor(), # 像素值从 0-255 归一化到 0-1 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) ) val_dataset = datasets.ImageFolder( root='data/val', 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]) ]) ) # 验证一下类别和索引的映射对不对 print(train_dataset.class_to_idx) # 输出 {'cat': 0, 'dog': 1}

这里的数据增强是一个容易被低估的环节。训练集加了RandomHorizontalFlip，也就是把图片随机左右翻转，这一步能显著降低模型过拟合。猫狗的轮廓在水平翻转后仍然是正常的猫狗，这是这类生物图像任务里最安全的数据增强方式。

Normalize里的 mean 和 std 用的是 ImageNet 数据集的统计值，这是行业惯例。原因很简单：后面要加载的预训练权重是基于 ImageNet 训练的，输入数据的分布最好和预训练时一致，否则迁移效果会打折扣。

3.2 训练集和验证集划分：别把数据洗牌这个环节省掉

很多新手拿到数据集后直接把所有图片塞进训练，最后模型在训练集上准确率 98%，一到验证就崩，这就是典型的没有划分数据。分类任务里训练集和验证集的分布必须接近，否则验证结果没有参考意义。

import os, shutil, random from glob import glob # 把原始图片文件按 8:2 划分到 train 和 val 目录 random.seed(42) for cls in ['cat', 'dog']: images = glob(f'raw/{cls}/*.jpg') random.shuffle(images) train_count = int(len(images) * 0.8) for img in images[:train_count]: shutil.copy(img, f'data/train/{cls}/') for img in images[train_count:]: shutil.copy(img, f'data/val/{cls}/')

random.seed这行不能省。不固定随机种子的话，每次跑脚本划分结果都不一样，后面你会发现今天训练集准确率 96%，明天同样的代码只剩 90%，实际上就是数据划分变了。固定随机种子后，对比实验就公平了。

3.3 数据加载的批大小和线程数设置

DataLoader 的参数设置也有讲究。batch_size决定一次前向传播处理多少张图，num_workers决定用几个子进程去读图。

from torch.utils.data import DataLoader train_loader = DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True # 加速数据从内存到显存的传输 ) val_loader = DataLoader( val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True )

batch_size=32是一个比较稳的起点。显存不够就降到 16 或 8，显存够用可以升到 64，但提升不是线性的，在小数据集上大 batch 反而可能让模型收敛变慢。num_workers在 Windows 上经常会报错，因为 Windows 的 multiprocessing 行为和 Linux 不同，如果报错就把num_workers改成 0。

4. 模型训练与调优：从迁移学习到自定义训练循环

4.1 为什么不从零搭网络：预训练权重的价值远大于你的想象

猫狗分类这种任务，模型架构有很多选择：VGG、ResNet、EfficientNet、ViT 都能做。但我的建议很简单——上来就用 PyTorch 官方预训练好的模型，而不是自己从零写一个卷积网络。原因有两个：一是预训练模型已经在 ImageNet 上学过通用的图像特征，比如边缘、纹理、形状，这些特征在小数据集上迁移过来是巨大的优势；二是自己从零搭的网络容易在细节上出错，比如卷积核尺寸不对、全连接层维度接不上，排查起来很费时间。

import torch.nn as nn from torchvision import models # 加载预训练 ResNet18，这是小数据集上性价比最高的模型之一 model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) # 替换最后一层全连接层，改成二分类输出 model.fc = nn.Linear(model.fc.in_features, 2) # 把模型移动到 GPU（如果可用） device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device)

weights=models.ResNet18_Weights.IMAGENET1K_V1是 PyTorch 新版推荐的写法，老版本写pretrained=True会收到弃用警告。替换最后一层全连接层时，model.fc.in_features会自动读取原来的输入维度，所以不管你用的 ResNet 还是其他系列，这行代码都能自动适配。

4.2 损失函数和优化器选型：二分类问题的标准组合

二分类问题最标准的做法是输出维度设为 2，用CrossEntropyLoss作为损失函数。CrossEntropyLoss内部已经包含了 softmax 操作，所以模型最后一层不需要再手动加 softmax。

import torch.optim as optim criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001)

Adam 是自适应学习率的优化器，对于新手来说它比 SGD 好调。SGD 需要仔细调学习率、动量和权重衰减，一个没调好就收敛得很慢，而 Adam 的默认参数在大多数任务上都能得到一个还不错的起点。不过 Adam 也不是万能的，如果你的准确率到了后期怎么都提不上去，可以试试把优化器换成 SGD，学习率设 0.01，动量设 0.9。

4.3 训练循环：完整代码与每个参数的含义

这里把完整的训练代码写出来。这个训练循环是通用模板，换数据集、换模型、换任务，都能直接复用。

import torch from tqdm import tqdm num_epochs = 10 best_val_acc = 0.0 for epoch in range(num_epochs): # ===== 训练阶段 ===== model.train() running_loss = 0.0 correct = 0 total = 0 for images, labels in tqdm(train_loader): images, labels = images.to(device), labels.to(device) # 梯度清零，这步忘掉的话梯度会在参数更新时累加 optimizer.zero_grad() # 前向传播 + 计算损失 outputs = model(images) loss = criterion(outputs, labels) # 反向传播 + 参数更新 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() train_loss = running_loss / total train_acc = correct / total # ===== 验证阶段 ===== model.eval() val_correct = 0 val_total = 0 val_loss = 0.0 with torch.no_grad(): for images, labels in val_loader: 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) val_total += labels.size(0) val_correct += (predicted == labels).sum().item() val_acc = val_correct / val_total val_loss = val_loss / val_total print(f'Epoch [{epoch+1}/{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_val_acc: best_val_acc = val_acc torch.save(model.state_dict(), 'best_model.pth') print(f' -> 保存新最佳模型，Val Acc: {val_acc:.4f}')

这里有几个细节需要解释：

model.train()和model.eval()是很多人容易忽略的。model.train()会启用 BatchNorm 的统计更新和 Dropout，model.eval()则固定这些层的行为。如果不切换，训练完直接做验证，BatchNorm 的 running stats 还在更新，验证结果会偏乐观。

torch.no_grad()在验证阶段必须加。它告诉 PyTorch 不需要计算梯度，这样前向传播的显存占用和耗时都会明显下降。不加也能跑，但比较浪费资源，而且等到模型变大、batch 变大以后，会发现验证时显存不够，其实原因就是这个。

torch.save(model.state_dict(), 'best_model.pth')只保存模型的权重参数，不保存完整的模型结构。加载时需要先定义好模型结构，再load_state_dict把权重填进去。这是 PyTorch 推荐的保存方式，理由是权重文件更小、跨版本兼容性更好。

4.4 学习率策略和训练轮数的经验参考

训练轮数 10 轮是起步值，猫狗分类数据量在几千到两万张这个区间，10 轮基本上能看到明显收敛，但最优轮数可能要到 15-20 轮。有一个常见现象值得注意：训练准确率每隔几轮会有一个明显的跳跃，这往往是因为模型从一个局部最优跳到了另一个更好的局部最优，不要在前几轮看到准确率没涨就提前终止。

学习率的调整策略有一个简单有效的方案——在训练到一半的时候把学习率降为原来的 1/10。用 PyTorch 实现起来很简洁：

# 每 5 个 epoch 把学习率乘以 0.1 scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1) # 在每个 epoch 结束时调用 scheduler.step()

这个策略的直觉是：训练初期需要大学习率快速找到好的参数区域，训练后期需要小学习率精细搜索最优值。你也可以试试ReduceLROnPlateau，它在验证损失不再下降时自动降低学习率，比固定间隔更智能，但要多观察一个参数patience。

下面用一个表格总结我建议的初始参数和调整方向：

5. 训练中的五个高频翻车现场：现象、原因与解法

5.1 loss 一直是 nan 或直接变成负数

现象：训练到某个 epoch 后，loss 变成 nan，或者出现负数，然后准确率也跟着崩掉。

原因：最典型的是学习率过大导致梯度爆炸。Adam 虽然是自适应学习率，但初始值 0.001 在某些数据分布下仍然偏高；另一个可能原因是输入数据里有异常的像素值，比如图片里有全黑图，归一化后出现除零错误。

解决：先看输入数据的取值范围，打印一下images.min()和images.max()确认是不是 0 到 1。如果确认数据没问题，把学习率降到 0.0001 再试。还有一个办法是在优化器里加clip_grad_norm_：

# 在 loss.backward() 之后加这行，梯度范数超过 1.0 就裁剪掉 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

5.2 验证准确率一直卡在 50% 左右

现象：训练准确率在上涨，但验证准确率纹丝不动，大概在 50% 徘徊——这相当于随机猜，模型什么都没学到。

原因：最常见的是标签和数据没对齐。比如 ImageFolder 要求的目录结构里，cat文件夹在dog前面，类别索引是 cat=0、dog=1，但在某个环节标签被转了，比如数据增强里做了一些奇怪的变换，或者验证集和训练集用了不同的预处理。

解决：先做一个小实验——用验证集的前 16 张图走一遍前向传播，打印预测值和真值，人工核对。如果发现数据没问题，那就把模型换回预训练权重，重新训练前先冻结特征提取层、只训练最后一个全连接层，跑 2-3 个 epoch 确认能收敛，再解冻全部参数。

# 冻结所有层的参数，只训练最后一层 for param in model.parameters(): param.requires_grad = False for param in model.fc.parameters(): param.requires_grad = True

5.3 CPU 训练速度慢到让人怀疑人生

现象：一个 epoch 跑了快半小时，训练集只有 2000 张图。

原因：除了确实没有 GPU 外，很可能是num_workers设成了 0，数据加载和模型计算串行执行，GPU（如果有的话）在等数据传到显存。Windows 平台上num_workers稍大一点还可能触发多进程报错。

解决：有 GPU 就确认一下程序是否真的用上了 GPU——在训练循环里加一行assert next(model.parameters()).is_cuda。如果是纯 CPU 跑，就把num_workers在 Linux 上调到 4、Windows 上调到 2，同时把pin_memory打开。数据加载是 CPU 训练最容易忽略的瓶颈，很多时候调好这两个参数，速度能快一倍。

5.4 GPU 显存不足，但这其实不是硬件问题

现象：训练到中途报CUDA out of memory，但nvidia-smi看显存占用率并不高。

原因：这个坑比较隐蔽——PyTorch 的显存分配策略是：即使模型很小，也可能因为某个 batch 里图片尺寸不一致导致临时张量异常大。最常见的是数据没有统一 resize，或者 DataLoader 里的collate_fn拼接时出了问题。

解决：检查数据增强里Resize((224, 224))是否真的生效，采样打印几个images.shape看是不是都是[3, 224, 224]。另外，如果用了梯度累积，注意梯度在显存里是持续占用的，不放宽batch_size的情况下可以调小输入尺寸，比如Resize((160, 160))，显存占用会成平方地下降。

5.5 加载模型权重时报 shape 不匹配

现象：用torch.load加载别人给的模型，报size mismatch for fc.weight: expected size [2, 512], got [1000, 512]。

原因：这个报错几乎都出现在没有理解state_dict的映射关系。预训练模型最后一层输出是 1000 类，因为 ImageNet 有 1000 个类别；你改成 2 类输出后，fc.weight的维度就是[2, 512]，加载原来的权重当然不匹配。

解决：如果你自己训练好模型再加载，报这个错说明模型定义没对上，重建模型后再load_state_dict：

# 加载训练好的权重，必须先重建同样的模型结构 model = models.resnet18() model.fc = nn.Linear(model.fc.in_features, 2) model.load_state_dict(torch.load('best_model.pth', map_location='cpu')) model.to(device)

map_location='cpu'是从 GPU 训练的模型拿到 CPU 上推理时必加的，不加的话没有 GPU 的机器会直接报错找不到 CUDA 设备。

6. 模型推理与验证结果：把训练好的模型用起来的两条经验

训练完成只是第一步，真正让这个项目有价值的是把模型导出、部署到实际环境里。我的习惯是训练完立刻用一个小脚本做一次端到端推理，确认从图片到预测结果的整条链路是通的。

from PIL import Image import torchvision.transforms as transforms def predict_image(image_path, model, 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') tensor = transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): output = model(tensor) # softmax 把输出转成概率分布 prob = torch.softmax(output, dim=1) _, predicted = torch.max(output, 1) # idx_to_class 是从训练时保存的映射关系，别弄丢 return predicted.item(), prob.max().item() # 调用示例 # model = models.resnet18() # model.fc = nn.Linear(512, 2) # model.load_state_dict(torch.load('best_model.pth', map_location='cpu')) # pred, prob = predict_image('test_dog.jpg', model, torch.device('cpu')) # print(f'预测类别: {pred}, 置信度: {prob:.4f}')

这里有一个我反复强调的细节：推理阶段的预处理必须和训练阶段完全一致。很多人训练时加了RandomHorizontalFlip，推理时不加，这没问题；但 Resize 的尺寸、Normalize 的均值和标准差如果变了，模型输出的概率分布会立刻偏离训练时的状态。最稳妥的做法是在训练脚本里把 transform 定义成函数，推理脚本直接 import 用同一个函数。

最后一个进阶技巧：训练完的模型不只是用于图片单张推理的。在后面做真实项目时，可以把模型转成 TorchScript 格式，这样不依赖 Python 环境也能在 C++ 或者移动端部署。model.eval()之后torch.jit.script(model)，得到的是一个打包好的静态计算图，运行速度通常比 Python 推理快不少。

这个项目做到最后，我最深的体会是：深度学习的代码实现只是表层功夫，真正决定模型效果的是数据质量、预处理一致性、训练策略的选择，以及遇到问题时的排查路径。猫狗分类做完以后，我已经换过三次工业数据集任务，但每次回到这个最小项目上重新审视，都还能找到之前忽略的优化空间。希望这篇笔记帮你在 PyTorch 方向上少踩几个坑，顺顺利利跑通属于你自己的第一个图像分类项目。

本文还有配套的精品资源，点击获取