PyTorch实现猫狗识别：从数据预处理到迁移学习全流程

简介这是一份基于Python机器学习的猫狗识别分类完整项目面向高校人工智能、计算机等相关专业学生可用于毕业设计、课程设计或深度学习入门实践。项目整合了数据预处理、模型训练与测试等环节覆盖CNN与Swin Transformer两种网络结构并附带PyTorch实现代码、说明文档及论文能够帮助理解图像分类任务的完整流程。压缩包共十六个文件约1.67MB以.py源码为主辅以.md说明、.txt数据、.docx论文和.pth模型权重文件目录清晰、便于按需查阅目前已有68人学习下载。资源内含训练好的CNN模型权重epoch400及完整训练日志可直接加载模型进行推理也可基于现有代码修改网络结构或参数灵活用于其他图像分类场景。对于基础薄弱的读者文档中还提供了配置与运行说明适合边学边练、快速上手。1. 先想清楚再动手猫狗识别到底在解什么问题猫狗识别分类是机器学习里最经典的图像二分类任务几乎每个入门者都会拿它练手但真正把它做成一个可用项目而不是一个跑完就扔的 Notebook中间隔着一整套工程决策。数据怎么组织、模型选多大、训练多久、验证集怎么切每一步都影响最终交付的模型能不能拿得出手。本文会从数据集预处理、模型选型、训练参数到验证推理把一条完整可复现的路径讲清楚适合已经会用 Python 但还没系统做过图像分类项目的人也适合想把自己手里的分类任务整理成规范工程的开发者。这个任务表面上是“识别猫和狗”实际上要解决的是图像分类的标准流程问题图片尺寸不一致、数据量有限、训练过程过拟合、模型文件如何保存与加载。把这些问题一个个解决掉你得到的不仅是一个猫狗分类器而是一套可以平移到其他二分类或小规模多分类任务的模板。下面直接从数据准备开始这是整个项目里最容易出错也最影响结果的一步。2. 准备数据集与预处理脏数据比模型更影响精度2.1 数据集选型公开数据集怎么下、怎么组织目录猫狗分类最常用的公开数据集是 Kaggle 的 Dogs vs. Cats包含 25000 张猫狗图片训练集 20000 张测试集 5000 张文件名如cat.100.jpg和dog.200.jpg。如果你不方便从 Kaggle 下载常见的替代方案是小鱼数据集或自行爬取但爬取的数据必须人工清洗否则标签噪声会让模型精度上限大打折扣。拿到数据后第一件事不是写模型而是把文件按下面的目录结构组织好这会直接决定后面代码的简洁程度data/ ├── train/ │ ├── cat/ # 猫图片 │ └── dog/ # 狗图片 ├── val/ │ ├── cat/ │ └── dog/ └── test/ # 无标签图片用于最终预测从文件名解析标签并移动到对应目录用 Python 标准库就能完成不需要引入额外依赖。按 8:2 比例随机切分训练集和验证集切分时要保证随机种子固定否则每次运行得到的数据划分不同实验结果无法复现。2.1.1 从文件名生成目录的脚本import os import shutil import random from pathlib import Path random.seed(42) source_dir Path(raw_data) train_dir Path(data/train) val_dir Path(data/val) for path in source_dir.iterdir(): if not path.is_file(): continue if path.name.startswith(cat): label cat elif path.name.startswith(dog): label dog else: continue dest_dir train_dir / label if random.random() 0.2 else val_dir / label dest_dir.mkdir(parentsTrue, exist_okTrue) shutil.copy(path, dest_dir / path.name)这个脚本把源目录里的图片按文件名前缀分到 cat 和 dog 两个类别下并随机抽出 20% 作为验证集。随机种子固定为 42保证每次执行结果一致。shutil.copy是复制而不是移动避免误删原始数据这在数据量不大时更安全。2.2 预处理管线Resize、归一化与数据增强的默认配置图片预处理直接决定模型能否收敛。原始数据集里的图片尺寸参差不齐有的横向、有的纵向、有的只有几十 KB而 PyTorch 的卷积网络要求输入张量形状固定。最常见的做法是把所有图片统一缩放到 224x224因为 ImageNet 预训练模型的标准输入尺寸就是 224。归一化必须使用预训练模型对应的均值和标准差。如果用 ResNet18 的 ImageNet 权重那么均值[0.485, 0.456, 0.406]、标准差[0.229, 0.224, 0.225]是固定的不能随意改。改了之后预训练权重的分布就被破坏模型需要重新学习低级特征迁移学习的效果会明显下降。数据增强只在训练集上用验证集和测试集只用 Resize 和归一化。这是很多初学者容易忽略的验证集如果也做了随机裁剪和翻转评估结果就不稳定甚至会把随机的增强结果当作模型的真实表现。训练集增强的标准配置是随机水平翻转、随机旋转和颜色抖动这能在不增加数据量的情况下扩大数据分布覆盖。2.2.1 PyTorch 的 transforms 配置from torchvision import transforms train_transforms transforms.Compose([ transforms.Resize(256), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), 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(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])先Resize(256)再RandomResizedCrop(224)是训练集的标准做法RandomResizedCrop会在缩放的同时随机裁剪不同区域相当于做了尺度增强。验证集用CenterCrop固定裁剪中心区域保证每次评估的输入一致。注意ToTensor必须在归一化之前因为它会把 HWC 的 PIL 图片转成 CHW 的 Tensor并且把像素值从 0~255 缩放到 0~1。2.3 用 ImageFolder 加载数据代码里不写死路径torchvision.datasets.ImageFolder是处理按类别分目录的图片数据最省事的工具它自动扫描子目录名作为类别标签并且按字母序分配索引。以data/train/cat里的图片为例cat文件夹会被映射为索引 0dog映射为索引 1。这样训练代码里就不用自己写标签解析逻辑换数据集也不需要改代码。from torchvision import datasets from torch.utils.data import DataLoader train_dataset datasets.ImageFolder(rootdata/train, transformtrain_transforms) val_dataset datasets.ImageFolder(rootdata/val, transformval_transforms) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue)shuffleTrue只在训练集打开验证集保持顺序即可。num_workers4让数据加载使用 4 个并行进程避免 GPU 等数据在 Windows 上要注意多进程数据加载可能报错如果遇到就把num_workers设为 0。pin_memoryTrue可以加速 GPU 上的数据传输前提是你的数据在 CPU 内存里能被页面锁定。3. 模型骨架选 ResNet18把训练成本压到最低3.1 为什么不自己搭 CNN小数据集的泛化瓶颈很多人会尝试自己写几层卷积加池化再拼全连接层训练下来发现验证集准确率只有 70%~80%差距往往不在模型结构本身而在数据量。猫狗数据集只有两万张图自己搭的 CNN 从头训练底层卷积核很难学到足够通用的边缘、纹理特征容易过拟合到训练集的颜色和背景上。常见做法是用 ImageNet 上预训练好的模型做迁移学习。ImageNet 的 1000 类分类任务让模型学会了通用的图像特征提取能力这些底层特征对猫狗识别依然有效。在两万张图片的小数据集上微调一个预训练 ResNet18 通常能在验证集上达到 95% 以上而且只需要几十分钟的训练时间。3.2 迁移学习微调冻结 backbone 与全量微调的取舍用torchvision.models一行代码就能加载预训练模型关键是把最后一层全连接替换成输出 2 个类别的线性层。这里要决定哪些参数参与反向传播冻结 backbone 的方式是把requires_grad设为False只训练最后一层。这种方式收敛最快适合数据量很小、希望快速得到基线效果的场景。但缺点是模型顶层的特征是根据 ImageNet 的任务学出来的在猫狗这种细粒度相差不是特别大的任务上效果尚可如果换到更独特的域比如 X 光图像冻结 backbone 就会成为瓶颈。全量微调则是让所有参数都参与训练用较小的学习率微调整个网络。这个方案训练时间更长但精度上限更高。我一般会先用冻结方式跑一轮拿到基线再解冻全部参数用小学习率微调。这种方式比直接全量训练更容易收敛因为前几轮已经让分类头适应了当前数据分布。import torchvision.models as models import torch.nn as nn model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) for param in model.parameters(): param.requires_grad False num_features model.fc.in_features model.fc nn.Linear(num_features, 2)第一轮训练时conv1到layer4的所有参数都被冻结只有新的fc层会更新。model.fc.in_features读取原全连接层的输入维度ResNet18 是 512替换成输出 2 的分类头。如果后续想要全量微调只需要把冻结参数那段去掉重新遍历一遍model.parameters()把requires_grad设回True即可。3.3 损失函数、优化器与学习率的初始化参数二分类任务最直接的损失函数是CrossEntropyLoss它内部集成了 softmax不需要在模型最后一层额外加激活函数。如果输出层是 2 维的 logits传入CrossEntropyLoss会自动完成 softmax 并计算交叉熵。优化器方面冻结 backbone 时用 Adam 足够学习率 1e-3 起步比较稳全量微调时用 SGD momentum 效果更稳妥学习率要降一个量级到 1e-4。这里有个值得注意的细节新加的fc层是随机初始化的而预训练部分的参数已经是收敛状态两者的梯度尺度差异大统一用同一个学习率时新的分类头学得太慢。处理方式是把参数分组对fc层用较大的学习率对特征提取层用较小的学习率param_groups [ {params: [p for n, p in model.named_parameters() if fc not in n], lr: 1e-4}, {params: [p for n, p in model.named_parameters() if fc in n], lr: 1e-3}, ] optimizer torch.optim.Adam(param_groups)named_parameters返回参数名和参数本身通过名字判断是否属于fc层。这样新分类头以 10 倍学习率快速适配到当前数据而预训练特征层只在局部范围内微调避免破坏已有特征。4. 训练循环与早停策略让验证集真正帮你做决策4.1 最小可用的 PyTorch 训练循环训练循环是每个 PyTorch 项目的核心骨架。这里给出一个直接可用的版本包含单 epoch 训练、验证、模型保存三个部分用 Tqdm 显示进度条方便观察训练状态。import torch import torch.nn as nn from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() def train_one_epoch(model, loader, optimizer, criterion, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in tqdm(loader, descTraining): 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() _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss running_loss / len(loader) accuracy correct / total return avg_loss, accuracyoptimizer.zero_grad()必须在每次反向传播前清空梯度PyTorch 的梯度是累加的不清空会把上一个 batch 的梯度叠加到当前 batch 上。torch.max(outputs, 1)返回每行最大值的值和索引索引就是预测的类别。model.train()把模型切到训练模式这会启用 Dropout 和 BatchNorm 的训练行为验证时必须切回model.eval()。4.2 早停与模型保存读验证集准确率而不是训练 loss训练时只看训练集 loss 是常见误区训练 loss 会一直下降但这不代表模型泛化能力在变好。要盯着验证集准确率当验证集准确率连续 N 个 epoch 不提升时就应该停止训练并恢复最优权重。best_acc 0.0 patience 5 bad_epochs 0 for epoch in range(10): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(model, val_loader, criterion, device) print(fEpoch {epoch}: train_acc{train_acc:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) bad_epochs 0 else: bad_epochs 1 if bad_epochs patience: print(Early stopping) breaktorch.save只保存state_dict而不是整个模型对象。state_dict保存的是所有层的权重和偏置体积更小加载时也只需要重建模型结构再载入权重对环境依赖更少。patience设为 5 意思是允许模型在验证集上连续 5 个 epoch 没有进步超过就停止这是控制过拟合最直接的手段。4.2.1 加载最优模型做推理model models.resnet18(weightsNone) model.fc nn.Linear(512, 2) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.to(device) model.eval()加载时必须先定义和保存时完全一致的模型结构。weightsNone表示不加载预训练权重因为我们要从本地文件恢复训练好的参数。map_locationdevice在只有 CPU 的机器上加载 GPU 训练好的权重时特别重要它能自动把张量从 cuda 映射到 cpu省去手动改设备的问题。4.3 常见坑数据加载慢、显存不足、准确率震荡数据加载慢时先看num_workers如果 CPU 核心数多可以把num_workers设成 CPU 核心数的一半。显存不足时把batch_size从 32 降到 16 或 8而不是调小图片尺寸因为图片尺寸会影响模型输入分布减弱迁移学习的效果。还有一个容易忽略的处理如果训练时没关其他占用显存的程序显存不足往往不是模型问题而是环境问题。准确率震荡不下降时先用小数据跑一批看看能不能过拟合。比如只拿 100 张图片训练 20 个 epoch如果训练准确率能到接近 100%说明数据和模型没问题是学习率或数据增强强度的问题。如果连小数据都过拟合不了检查数据预处理有没有把标签和图片对应错ImageFolder 常见的坑是子目录顺序不对导致标签翻转。5. 验证模型与批量预测脚本把模型交到别人手里做完训练项目交付的核心是一份能单张预测的脚本。这里提供一个直接从命令行接收图片路径的predict.py脚本用argparse解析参数输出类别和置信度。模型结构要和训练时保持一致同时复用验证集的预处理管线。import argparse import torch from PIL import Image from torchvision import transforms def predict_image(model, image_path, device): transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(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) with torch.no_grad(): outputs model(tensor) probabilities torch.softmax(outputs, dim1) confidence, predicted torch.max(probabilities, 1) return predicted.item(), confidence.item() if __name__ __main__: device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsNone) model.fc nn.Linear(512, 2) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.to(device) model.eval() parser argparse.ArgumentParser() parser.add_argument(--image, typestr, requiredTrue) args parser.parse_args() class_names [cat, dog] idx, conf predict_image(model, args.image, device) print(fPredicted: {class_names[idx]}, confidence: {conf:.4f}).convert(RGB)必须加因为有些图片是灰度图或带透明通道的 RGBA 图直接训练好的模型要求三通道输入不转换会报形状不匹配。torch.no_grad()在推理时禁用梯度计算能减少显存占用并提升推理速度。最后推荐的验证技巧是把训练时没有见过的图片单独放一个目录做批量预测统计每个类别的平均置信度——如果模型对某一类普遍低置信度说明那一类样本在训练集中占比不足需要补充数据或增加数据增强强度这比单一准确率更能反映模型在生产环境下的可靠性。本文还有配套的精品资源点击获取