简介这是一份面向深度学习初学者与有一定基础的开发者的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()就返回 FalseGPU 完全没被识别到。这个问题的根源几乎都是 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 python3.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 上可能要十几分钟但如果只是验证代码逻辑、跑通训练流程用一个小数据集或者只训练几个 epochCPU 完全够用。我自己习惯的做法是先在一小部分数据上把完整流程跑通再切到 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( rootdata/train, transformtransforms.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( rootdata/val, transformtransforms.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(fraw/{cls}/*.jpg) random.shuffle(images) train_count int(len(images) * 0.8) for img in images[:train_count]: shutil.copy(img, fdata/train/{cls}/) for img in images[train_count:]: shutil.copy(img, fdata/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_size32, shuffleTrue, num_workers4, pin_memoryTrue # 加速数据从内存到显存的传输 ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue )batch_size32是一个比较稳的起点。显存不够就降到 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(weightsmodels.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)weightsmodels.ResNet18_Weights.IMAGENET1K_V1是 PyTorch 新版推荐的写法老版本写pretrainedTrue会收到弃用警告。替换最后一层全连接层时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(), lr0.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(fEpoch [{epoch1}/{num_epochs}], fTrain Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, fVal 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 的统计更新和 Dropoutmodel.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_size5, gamma0.1) # 在每个 epoch 结束时调用 scheduler.step()这个策略的直觉是训练初期需要大学习率快速找到好的参数区域训练后期需要小学习率精细搜索最优值。你也可以试试ReduceLROnPlateau它在验证损失不再下降时自动降低学习率比固定间隔更智能但要多观察一个参数patience。下面用一个表格总结我建议的初始参数和调整方向参数初始值常见调整方向batch_size32显存不够降到 16准确率波动大可以尝试 64学习率0.001Adam不收敛就降到 0.0001收敛太慢升到 0.002优化器Adam后期精度不够换 SGD momentum训练轮数10看验证集是否还在上升可延长到 20输入尺寸224x224小图可用 128x128 加速大图用 256x2565. 训练中的五个高频翻车现场现象、原因与解法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_norm1.0)5.2 验证准确率一直卡在 50% 左右现象训练准确率在上涨但验证准确率纹丝不动大概在 50% 徘徊——这相当于随机猜模型什么都没学到。原因最常见的是标签和数据没对齐。比如 ImageFolder 要求的目录结构里cat文件夹在dog前面类别索引是 cat0、dog1但在某个环节标签被转了比如数据增强里做了一些奇怪的变换或者验证集和训练集用了不同的预处理。解决先做一个小实验——用验证集的前 16 张图走一遍前向传播打印预测值和真值人工核对。如果发现数据没问题那就把模型换回预训练权重重新训练前先冻结特征提取层、只训练最后一个全连接层跑 2-3 个 epoch 确认能收敛再解冻全部参数。# 冻结所有层的参数只训练最后一层 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True5.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_locationcpu)) model.to(device)map_locationcpu是从 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, dim1) _, 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_locationcpu)) # 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 方向上少踩几个坑顺顺利利跑通属于你自己的第一个图像分类项目。本文还有配套的精品资源点击获取
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()就返回 FalseGPU 完全没被识别到。这个问题的根源几乎都是 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 python3.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 上可能要十几分钟但如果只是验证代码逻辑、跑通训练流程用一个小数据集或者只训练几个 epochCPU 完全够用。我自己习惯的做法是先在一小部分数据上把完整流程跑通再切到 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( rootdata/train, transformtransforms.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( rootdata/val, transformtransforms.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(fraw/{cls}/*.jpg) random.shuffle(images) train_count int(len(images) * 0.8) for img in images[:train_count]: shutil.copy(img, fdata/train/{cls}/) for img in images[train_count:]: shutil.copy(img, fdata/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_size32, shuffleTrue, num_workers4, pin_memoryTrue # 加速数据从内存到显存的传输 ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue )batch_size32是一个比较稳的起点。显存不够就降到 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(weightsmodels.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)weightsmodels.ResNet18_Weights.IMAGENET1K_V1是 PyTorch 新版推荐的写法老版本写pretrainedTrue会收到弃用警告。替换最后一层全连接层时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(), lr0.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(fEpoch [{epoch1}/{num_epochs}], fTrain Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, fVal 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 的统计更新和 Dropoutmodel.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_size5, gamma0.1) # 在每个 epoch 结束时调用 scheduler.step()这个策略的直觉是训练初期需要大学习率快速找到好的参数区域训练后期需要小学习率精细搜索最优值。你也可以试试ReduceLROnPlateau它在验证损失不再下降时自动降低学习率比固定间隔更智能但要多观察一个参数patience。下面用一个表格总结我建议的初始参数和调整方向参数初始值常见调整方向batch_size32显存不够降到 16准确率波动大可以尝试 64学习率0.001Adam不收敛就降到 0.0001收敛太慢升到 0.002优化器Adam后期精度不够换 SGD momentum训练轮数10看验证集是否还在上升可延长到 20输入尺寸224x224小图可用 128x128 加速大图用 256x2565. 训练中的五个高频翻车现场现象、原因与解法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_norm1.0)5.2 验证准确率一直卡在 50% 左右现象训练准确率在上涨但验证准确率纹丝不动大概在 50% 徘徊——这相当于随机猜模型什么都没学到。原因最常见的是标签和数据没对齐。比如 ImageFolder 要求的目录结构里cat文件夹在dog前面类别索引是 cat0、dog1但在某个环节标签被转了比如数据增强里做了一些奇怪的变换或者验证集和训练集用了不同的预处理。解决先做一个小实验——用验证集的前 16 张图走一遍前向传播打印预测值和真值人工核对。如果发现数据没问题那就把模型换回预训练权重重新训练前先冻结特征提取层、只训练最后一个全连接层跑 2-3 个 epoch 确认能收敛再解冻全部参数。# 冻结所有层的参数只训练最后一层 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True5.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_locationcpu)) model.to(device)map_locationcpu是从 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, dim1) _, 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_locationcpu)) # 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 方向上少踩几个坑顺顺利利跑通属于你自己的第一个图像分类项目。本文还有配套的精品资源点击获取
