☰
PyTorch猫狗图像分类实战:从数据管道到ONNX部署全链路
2026/10/6 15:05:17 网站建设 项目流程

简介:这份资源是面向深度学习初学者与希望深入掌握 PyTorch 的开发者的一份猫狗图像分类实战教程,以 docx 文档形式呈现,聚焦卷积神经网络与数据增强等核心技能,帮助读者从零走通图像分类项目的完整流程。压缩包内共 1 个 docx 文件,约 19KB,内容涵盖项目背景、数据准备与增强、轻量级 CNN 模型构建、训练评估及部署步骤,并配有可直接复制运行的 Python 代码片段与性能调优技巧。教程还总结了常见陷阱与改进方向,既适合个人自学快速上手,也可作为教学辅助材料或技术博客的写作模板。目前已有 152 人学习,适合需要一份条理清晰、案例与代码兼备的入门实战参考的读者。

1. 从一堆毛茸茸的照片说起:这套 PyTorch 猫狗分类资源到底能跑出什么

你手头有 25000 张猫狗照片,想训练一个能区分它们的模型,但打开 PyTorch 官方文档发现示例全是 MNIST 手写数字,换成自己的数据就报错——这不是你一个人的问题。这套基于 PyTorch 的猫狗图像分类实战项目,核心就是解决“从标准数据集到真实图片”的断层。它提供了一条完整的可复现路径:用ImageFolder加载自定义目录结构,用轻量级 CNN 做二分类,配合数据增强和迁移学习把验证集准确率推到 90% 以上。适合刚学完卷积神经网络理论、想动手跑通第一个真实图像分类任务的开发者,也适合需要一份结构清晰的教学案例的讲师。资源本身不依赖特殊硬件,一张 6GB 显存的显卡就能跑,CPU 训练虽然慢但也能出结果。下面我从数据组织、模型搭建、训练调参到部署推理,把这条链路拆开讲透,重点放在那些教程里不写、但一跑就报错的地方。

2. 数据管道与增强策略:把 25000 张图喂进模型之前要做的四件事

2.1 目录结构决定你能不能跑通第一行代码

PyTorch 的ImageFolder对目录层级有硬性要求:每个类别一个文件夹,文件夹名就是标签。很多人把 Kaggle 下载的PetImages直接丢进去,结果发现里面混着Cat和Dog两个子目录,但还有Cat/666.jpg这种损坏文件,一加载就抛UnidentifiedImageError。常见做法是先写一个清洗脚本,把尺寸为 0 的图片和无法打开的图片剔除。

import os from PIL import Image def clean_dataset(root_dir): """遍历目录,删除损坏或空图片""" removed = 0 for subdir, _, files in os.walk(root_dir): for fname in files: fpath = os.path.join(subdir, fname) try: img = Image.open(fpath) img.verify() # 验证文件完整性 if os.path.getsize(fpath) == 0: raise ValueError("empty file") except Exception: os.remove(fpath) removed += 1 print(f"清理完成,删除 {removed} 个损坏文件") clean_dataset("PetImages")

这段代码做了两件事:img.verify()检查图片是否可解码,getsize排除空文件。参数root_dir指向你的数据集根目录,脚本会递归处理所有子文件夹。跑完后再按 8:2 划分训练集和验证集,建议用splitfolders库或者手动shutil.move,保证每个类别在验证集中都有足够样本。

2.2 训练集和验证集的增强策略必须分开写

数据增强是提升泛化能力最便宜的手段,但很多人把同一套transform同时用在训练和验证上,导致验证指标虚高。正确做法是:训练集用随机裁剪、翻转、颜色抖动,验证集只做缩放和中心裁剪。

import torchvision.transforms as T # 训练集:带随机增强 train_tf = T.Compose([ T.RandomResizedCrop(224, scale=(0.7, 1.0)), T.RandomHorizontalFlip(p=0.5), T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 验证集:只做确定性预处理 val_tf = T.Compose([ T.Resize(256), T.CenterCrop(224), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])

RandomResizedCrop(224, scale=(0.7, 1.0))表示随机裁剪出原图 70% 到 100% 的区域再缩放到 224,这个参数别设太小,否则猫耳朵狗尾巴容易被裁掉。ColorJitter的四个参数控制亮度、对比度、饱和度、色调的扰动幅度,0.2 是个保守值,再大可能让毛色失真。归一化用的均值和标准差是 ImageNet 统计值,这是迁移学习的通用做法,即使你的数据集不是 ImageNet,用这组参数也不会出问题。

2.3 DataLoader 的 num_workers 和 pin_memory 怎么设

DataLoader有两个参数直接影响训练速度:num_workers和pin_memory。在 Windows 上num_workers设大于 0 可能报BrokenPipeError,这是多进程启动方式的问题,常见做法是加if __name__ == "__main__":保护或者直接设 0。Linux 下可以设成 CPU 核心数的一半。

from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_ds = ImageFolder("train_data", transform=train_tf) val_ds = ImageFolder("val_data", transform=val_tf) train_loader = DataLoader( train_ds, batch_size=32, shuffle=True, num_workers=4, pin_memory=True ) val_loader = DataLoader( val_ds, batch_size=32, shuffle=False, num_workers=4, pin_memory=True )

pin_memory=True会把数据锁在内存页中,加速 CPU 到 GPU 的传输,显存充足时建议开启。batch_size=32是 6GB 显存下的安全值,如果你用 ResNet18 且图片尺寸 224,可以试 64,但要注意CUDA out of memory报错时先降 batch size 而不是改模型。

3. 模型搭建与迁移学习:从三层 CNN 到 ResNet18 的准确率跃迁

3.1 手写 CNN 的通道数和全连接层尺寸怎么算

项目正文给了一个两层卷积加两层全连接的CatDogClassifier,结构清晰但有个容易翻车的地方:全连接层的输入维度32*56*56是硬编码的,一旦你改了输入图片尺寸或者卷积层配置,这个数字就对不上。正确做法是用torch.flatten或者动态计算。

import torch import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes=2): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(2), # 224 -> 112 nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), # 112 -> 56 nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)) # 自适应池化,输出 64x1x1 ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(64, 128), nn.ReLU(), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): x = self.features(x) return self.classifier(x)

关键改动是AdaptiveAvgPool2d((1, 1)),它把任意尺寸的特征图压成 1x1,这样全连接层输入固定为通道数 64,不用再手算32*56*56。BatchNorm2d加在卷积和 ReLU 之间,能加速收敛并允许更大的学习率。Dropout(0.5)放在全连接层之间,防止过拟合。

3.2 迁移学习为什么比从头训练更划算

猫狗数据集只有 25000 张图,从头训练一个深层网络很容易过拟合。ResNet18 在 ImageNet 上预训练过的权重已经学会了边缘、纹理、形状等通用特征,你只需要替换最后的全连接层,用较小的学习率微调。

import torchvision.models as models def build_resnet18(num_classes=2, freeze_backbone=True): model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) if freeze_backbone: for param in model.parameters(): param.requires_grad = False # 替换最后的全连接层 in_features = model.fc.in_features model.fc = nn.Sequential( nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, num_classes) ) return model

freeze_backbone=True会冻结除fc层以外的所有参数,只训练新加的分类头。这样训练速度快,显存占用低,适合数据量小的场景。如果验证集准确率卡在 85% 上不去,可以把freeze_backbone设为False,用lr=1e-4微调整个网络,通常能再涨 3 到 5 个百分点。注意weights参数在新版 torchvision 里替代了旧的pretrained=True,旧写法会报警告。

3.3 损失函数和优化器的选择逻辑

二分类问题可以用CrossEntropyLoss也可以用BCELoss。CrossEntropyLoss配合 2 维输出更通用,扩展多分类时不用改代码。优化器首选Adam,学习率 1e-3 起步;如果微调整个 ResNet,换成SGD加动量 0.9 和权重衰减 1e-4 往往效果更好。

import torch.optim as optim device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = build_resnet18(freeze_backbone=True).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3, weight_decay=1e-4 ) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

filter(lambda p: p.requires_grad, ...)只把需要梯度的参数传给优化器,冻结的层不更新。StepLR每 5 个 epoch 把学习率乘 0.5,帮助模型在后期稳定收敛。weight_decay=1e-4是 L2 正则化,抑制过拟合。

4. 训练循环与验证:早停、学习率调度和显存监控的实操细节

4.1 训练循环里必须记录哪些指标

一个完整的训练循环要记录训练损失、训练准确率、验证损失、验证准确率。只看损失容易忽略过拟合,只看准确率可能掩盖类别不平衡。

def train_one_epoch(model, loader, criterion, optimizer, device): model.train() running_loss, correct, total = 0.0, 0, 0 for images, labels in 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) preds = outputs.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) return running_loss / total, correct / total @torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() running_loss, correct, total = 0.0, 0, 0 for images, labels in loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) running_loss += loss.item() * images.size(0) preds = outputs.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) return running_loss / total, correct / total

loss.item() * images.size(0)是加权平均,因为最后一个 batch 可能不满。@torch.no_grad()装饰器关闭验证阶段的梯度计算,省显存。model.eval()切换 BatchNorm 和 Dropout 到推理模式,这一步漏了会导致验证结果不稳定。

4.2 早停机制怎么写才不误杀

早停的逻辑是:验证损失连续 N 个 epoch 不下降就停止训练,并恢复验证损失最低时的权重。N 一般设 3 到 5,太小容易在震荡期误停,太大浪费训练时间。

best_val_loss = float("inf") patience, trigger = 5, 0 best_state = None for epoch in range(30): train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc = evaluate(model, val_loader, criterion, device) scheduler.step() print(f"Epoch {epoch+1}: train_loss={train_loss:.4f} train_acc={train_acc:.4f} " f"val_loss={val_loss:.4f} val_acc={val_acc:.4f}") if val_loss < best_val_loss: best_val_loss = val_loss best_state = model.state_dict().copy() trigger = 0 else: trigger += 1 if trigger >= patience: print(f"早停于 epoch {epoch+1},恢复最佳权重") model.load_state_dict(best_state) break

model.state_dict().copy()保存最佳权重,注意要.copy(),否则保存的是引用,后续训练会覆盖。scheduler.step()放在 epoch 末尾,不要放在 batch 循环里。

4.3 显存不够时先查这三个地方

CUDA out of memory是最常见的报错。排查顺序:第一,降batch_size,从 32 降到 16 甚至 8;第二,检查是否在验证阶段忘了torch.no_grad(),导致梯度图累积;第三,用torch.cuda.empty_cache()清理缓存,但这不是根本解决办法。如果还不行,把图片尺寸从 224 降到 128,显存占用会降到约三分之一。

import torch print(torch.cuda.memory_allocated() / 1024**2, "MB") print(torch.cuda.max_memory_allocated() / 1024**2, "MB")

这两行打印当前显存和峰值显存,帮你判断是模型太大还是中间变量没释放。

5. 避坑与排查:训练猫狗分类器时最容易翻车的五个地方

5.1 验证集准确率比训练集还高

现象:训练到第 3 个 epoch,验证准确率 92%,训练准确率只有 85%。原因:训练集用了Dropout和RandomResizedCrop等强增强,验证集只做中心裁剪,模型在验证集上“更容易”。另外BatchNorm在训练模式用 batch 统计量,验证模式用滑动平均,也会造成差异。解决:如果验证准确率持续高于训练准确率且差距在 5% 以内,属于正常现象;如果差距超过 10%,检查验证集是否混入了训练集图片。

5.2 损失变成 NaN

现象:训练几个 batch 后 loss 突然变成nan。原因:学习率太大导致梯度爆炸,或者输入数据没有归一化,像素值在 0 到 255 之间。解决:先把学习率降到 1e-4,确认Normalize在ToTensor之后执行。如果用的是自定义归一化参数,检查标准差是否写成了 0。

5.3 预测结果全是猫或全是狗

现象:模型在测试集上把所有图片都预测成同一类。原因:类别不平衡,比如训练集里猫 12000 张、狗 8000 张,模型学到“全猜猫”就能拿到 60% 准确率。解决:用WeightedRandomSampler过采样少数类,或者在CrossEntropyLoss里传weight参数。

from torch.utils.data import WeightedRandomSampler targets = [label for _, label in train_ds.samples] class_counts = torch.bincount(torch.tensor(targets)) weights = 1.0 / class_counts.float() sample_weights = weights[torch.tensor(targets)] sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True) train_loader = DataLoader(train_ds, batch_size=32, sampler=sampler, num_workers=4)

5.4 加载模型推理时结果和训练时不一致

现象:训练时验证准确率 91%,保存模型后重新加载推理,准确率掉到 70%。原因:推理时忘了调用model.eval(),Dropout还在随机丢弃神经元,BatchNorm还在用 batch 统计量。解决:加载权重后立刻model.eval(),并用torch.no_grad()包裹推理代码。

5.5 Windows 上 num_workers 大于 0 报错

现象:RuntimeError: DataLoader worker (pid xxx) is killed by signal。原因:Windows 的多进程启动方式是spawn,不是fork,子进程会重新导入主模块,如果主模块没有if __name__ == "__main__":保护就会递归创建进程。解决:把训练代码包在if __name__ == "__main__":里,或者设num_workers=0。

6. 从 .pth 到 ONNX:模型导出与推理加速的一个关键技巧

训练完保存的.pth文件只能在 PyTorch 环境里加载,部署到其他框架或者移动端需要转成 ONNX。这一步有个容易忽略的细节:导出时的输入尺寸必须和推理时一致,动态轴要显式指定。

import torch model = build_resnet18(freeze_backbone=False) model.load_state_dict(torch.load("best_model.pth", map_location="cpu")) model.eval() dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, "cat_dog.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}, opset_version=11 )

dynamic_axes把 batch 维度设为动态,这样导出后的模型可以接受任意 batch size 的输入。opset_version=11兼容性较好,如果部署环境支持更高版本可以调到 13 或 17。导出后用onnxruntime验证一下输出是否和 PyTorch 一致。

import onnxruntime as ort import numpy as np sess = ort.InferenceSession("cat_dog.onnx") dummy = np.random.randn(1, 3, 224, 224).astype(np.float32) onnx_out = sess.run(None, {"input": dummy})[0] torch_out = model(torch.from_numpy(dummy)).detach().numpy() print("最大误差:", np.abs(onnx_out - torch_out).max())

误差在 1e-5 以内说明导出正确。如果误差很大,检查是否有自定义层不被 ONNX 支持,或者opset_version太低。

还有一个提速技巧:在导出 ONNX 之前把模型转成torch.jit.trace格式,能进一步优化计算图。但注意trace不支持动态控制流,如果你的模型里有if分支,要用torch.jit.script。

traced = torch.jit.trace(model, dummy_input) traced.save("cat_dog_traced.pt")

从那以后我每次训练完都会先跑一遍 ONNX 导出和误差验证,确认推理结果和训练时一致再部署。这个习惯帮我省掉了至少三次“线上准确率暴跌”的排查时间。希望帮到你。

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

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询