ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

PyTorch实战:猫狗图像分类从环境搭建到模型部署全流程

PyTorch实战:猫狗图像分类从环境搭建到模型部署全流程 简介这份资源是面向深度学习初学者与希望深入理解 PyTorch 的开发者打造的猫狗图像分类实战教程以 docx 文档形式呈现压缩包内共 1 个文件约 19KB。内容围绕图像分类全流程展开涵盖项目背景、数据准备与增强、轻量级 CNN 模型构建、训练与评估方法以及部署步骤并配有可直接复制运行的 Python 代码片段帮助读者在动手实践中掌握 PyTorch 求解实际问题的思路。教程还融入多种提升性能与准确率的训练技巧并在结尾总结常见陷阱与改进方向既适合个人学习快速入门也可作为教学辅助或技术博客写作的参考模板。目前已有 152 人学习适合需要一份条理清晰、案例与代码兼备的入门实战资料的读者。1. 猫狗分类项目为什么值得作为 PyTorch 实战第一站如果你刚配好 PyTorch 环境想找一个能跑通全流程、又不至于一上来就被劝退的项目猫狗图像分类几乎是公认的最佳起点。它同时踩中了深度学习入门最核心的几件事自定义数据集加载、卷积网络搭建、迁移学习微调、训练过程监控、模型保存与推理部署。数据量适中二分类任务收敛快单卡甚至 CPU 都能跑出结果但麻雀虽小五脏俱全工业界图像分类项目的骨架它一个不缺。很多人搜「深度学习实战项目案例」时真正卡住的不是模型写不出来而是数据怎么组织、训练不收敛怎么排查、验证集准确率虚高怎么办。这篇笔记就按一线做项目的顺序把基于 PyTorch 的猫狗图像分类从环境、数据、模型、训练到推理完整走一遍参数怎么设、坑在哪、为什么这么选都讲清楚。适合刚学完 CNN 理论想动手的人也适合想把这套流程迁移到自己业务数据集上的工程师。2. 环境搭建与数据准备把地基打牢再谈模型2.1 PyTorch 环境搭建的三种路径与选择环境这一步翻车的人比想象中多。搜「pytorch安装」「anaconda配置pytorch环境」「ubuntu 安装pytorch」的人十有八九是卡在 CUDA 版本和 PyTorch 版本对不上。先说结论如果你只是做猫狗分类这种规模的任务CPU 版本完全够用一张 2 万张的猫狗数据集CPU 训练一个 epoch 大概几分钟跑十几个 epoch 也能接受。有 NVIDIA 显卡再考虑 CUDA 版本。三种常见路径conda 创建独立环境最推荐隔离干净不会污染系统 Python。pip 直接安装轻量但依赖冲突时排查麻烦。WSL 里装Windows 用户想用 Linux 工具链又不想装双系统时的折中方案搜「pytorch环境搭建wsl」的人多半是这个诉求。conda 路径的具体命令# 创建独立环境Python 版本选 3.9 或 3.10兼容性最好 conda create -n catdog python3.10 -y conda activate catdog # 安装 PyTorchCPU 版本用这条 pip install torch torchvision torchaudio # 如果有 CUDA 11.8 显卡去官网查对应命令不要凭记忆写 # 典型形式是 pip install torch --index-url https://download.pytorch.org/whl/cu118安装完必须验证别装完就往下走import torch print(torch.__version__) # 确认版本号 print(torch.cuda.is_available()) # 有显卡应为 TrueCPU 版为 False print(torch.cuda.device_count()) # 显卡数量torch.cuda.is_available()返回 False 但你确实有显卡八成是 CUDA 版本和 PyTorch 编译版本不匹配或者驱动太旧。这时候别硬扛先卸干净重装pip uninstall torch torchvision再按官网命令装。血泪经验环境问题越早解决越省事带着一个半坏的环境往下做后面报错你根本分不清是代码问题还是环境问题。2.2 猫狗数据集的目录结构与 Dataset 写法猫狗数据集常见来源是 Kaggle 的 Dogs vs. Cats训练集各 12500 张。下载后通常是train/cat.0.jpg、train/dog.0.jpg这种命名。但直接拿原始目录训练有两个问题没有验证集标签靠文件名解析容易出错。我一般先做一次目录重组按ImageFolder的标准结构来data/ train/ cat/ cat.0001.jpg ... dog/ dog.0001.jpg ... val/ cat/ ... dog/ ...重组脚本import os, shutil, random src_dir train # 原始目录文件名形如 cat.0.jpg dst_dir data val_ratio 0.2 # 验证集比例 random.seed(42) # 固定随机种子保证可复现 for split in [train, val]: for cls in [cat, dog]: os.makedirs(os.path.join(dst_dir, split, cls), exist_okTrue) files [f for f in os.listdir(src_dir) if f.endswith(.jpg)] random.shuffle(files) for i, fname in enumerate(files): cls cat if fname.startswith(cat) else dog split val if i len(files) * val_ratio else train shutil.copy(os.path.join(src_dir, fname), os.path.join(dst_dir, split, cls, fname))这段逻辑很直白先建好四个目标目录打乱文件列表后按比例切分按文件名前缀判断类别。random.seed(42)是关键不固定种子的话每次切分结果不同实验没法对比。val_ratio一般取 0.1 到 0.2数据量小就取 0.2数据量大取 0.1 就够。重组完用ImageFolder加载它会自动按子目录名生成标签from torchvision import datasets, transforms train_tf transforms.Compose([ transforms.Resize((224, 224)), # 统一尺寸配合迁移学习 transforms.RandomHorizontalFlip(), # 随机水平翻转增强 transforms.RandomRotation(10), # 小角度旋转 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], # ImageNet 均值 [0.229, 0.224, 0.225]) # ImageNet 标准差 ]) val_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_ds datasets.ImageFolder(data/train, transformtrain_tf) val_ds datasets.ImageFolder(data/val, transformval_tf)这里有几个参数必须解释。Resize((224, 224))是因为后面要用 ResNet 这类在 ImageNet 上预训练的模型它们输入固定 224。Normalize的均值和标准差是 ImageNet 统计出来的用预训练权重就必须用这套值否则输入分布和预训练时不一致微调效果会打折。训练集加RandomHorizontalFlip和RandomRotation验证集不加这是铁律——验证集要模拟真实推理不能引入随机性。提示ImageFolder要求每个类别一个子目录目录名就是类别名。如果你的数据是cat.0.jpg这种平铺结构必须先重组别想着在 Dataset 里硬解析文件名后期维护成本高。3. 模型选型与迁移学习别从零训一个 CNN3.1 从零搭 CNN 还是用预训练模型搜「深度学习cnn」「图像分类算法」的人很多第一反应是自己搭一个卷积网络。能跑但猫狗这种任务从零训一个小 CNN 大概能到 85% 左右准确率而用预训练的 ResNet18 微调轻松上 97%。差距来自 ImageNet 上百万张图学到的通用特征——边缘、纹理、形状这些对猫狗分类同样有效。所以我的建议很明确默认用迁移学习。只有当你数据集和 ImageNet 差异极大比如医学影像、卫星图或者要做学术对比实验时才考虑从零训。搜「最新的图像分类模型」会看到 ConvNeXt、ViT 这些但作为实战第一站ResNet18 足够速度快、显存占用低、社区资料多出问题好查。选型对比模型参数量猫狗任务预期准确率训练速度适用场景自建小 CNN约 1M82%~86%快教学、理解原理ResNet18 预训练11M96%~98%中实战首选ResNet50 预训练25M97%~98.5%慢追求精度、显存够ViT-Base 预训练86M97%~99%很慢数据量大、有 GPU3.2 迁移学习的两种微调策略与代码迁移学习不是简单加载权重就完事关键在「冻哪些层、训哪些层」。两种常见策略策略一只训分类头。把预训练模型所有卷积层冻结只替换最后的全连接层并训练它。适合数据量小几千张、和目标域差异不大的情况。训练快不容易过拟合。策略二解冻部分层一起微调。冻结前面的底层特征解冻后面的高层和分类头一起训。适合数据量中等偏上、想榨取更高精度的情况。学习率要设小否则会把预训练学到的特征破坏掉。代码实现import torch.nn as nn from torchvision import models def build_model(strategyhead_only, num_classes2): # 加载预训练 ResNet18weights 参数用新版写法 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) # 先冻结所有参数 for p in model.parameters(): p.requires_grad False # 替换最后的全连接层这一层默认 requires_gradTrue in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) if strategy partial: # 解冻 layer4 和 fc让高层特征也参与微调 for p in model.layer4.parameters(): p.requires_grad True return model model build_model(strategypartial) # 只把需要梯度的参数交给优化器避免无谓计算 trainable [p for p in model.parameters() if p.requires_grad] print(f可训练参数量: {sum(p.numel() for p in trainable)})逻辑说明requires_grad False冻结参数反向传播不会更新它们。替换model.fc后新层的requires_grad默认是 True所以「只训分类头」策略下只有 fc 在训。partial策略额外解冻layer4这是 ResNet 的最后一个残差块负责高层语义特征微调它收益最大。参数说明weightsmodels.ResNet18_Weights.IMAGENET1K_V1是 torchvision 新版 API老代码里写pretrainedTrue新版本已废弃搜「pytorch适配」的人经常踩这个。num_classes2对应猫狗二分类如果你换成十分类任务改这个数就行其他不用动。注意优化器只传requires_gradTrue的参数。如果把冻结参数也传进去虽然不会更新但会浪费显存和计算数据量大时很明显。4. 训练循环与参数调优让模型真正收敛4.1 训练循环的标准骨架训练循环看着简单但每个环节都有讲究。下面是一个完整可用的骨架import torch from torch.utils.data import DataLoader from torch import optim from torch.optim.lr_scheduler import StepLR device torch.device(cuda if torch.cuda.is_available() else cpu) model build_model(strategypartial).to(device) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) criterion nn.CrossEntropyLoss() # 微调时学习率要小1e-4 是安全起点 optimizer optim.Adam(trainable, lr1e-4, weight_decay1e-4) scheduler StepLR(optimizer, step_size5, gamma0.5) def run_epoch(loader, trainTrue): model.train() if train else model.eval() total_loss, correct, total 0.0, 0, 0 torch.set_grad_enabled(train) # 验证时关闭梯度省显存 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) if train: optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) if train: loss.backward() optimizer.step() total_loss loss.item() * imgs.size(0) correct (outputs.argmax(1) labels).sum().item() total imgs.size(0) return total_loss / total, correct / total for epoch in range(15): tr_loss, tr_acc run_epoch(train_loader, trainTrue) va_loss, va_acc run_epoch(val_loader, trainFalse) scheduler.step() print(fEpoch {epoch1:02d} | train {tr_loss:.4f}/{tr_acc:.4f} f| val {va_loss:.4f}/{va_acc:.4f} | lr {optimizer.param_groups[0][lr]:.2e})逻辑说明model.train()和model.eval()切换模式影响 BatchNorm 和 Dropout 的行为忘了切会导致验证结果不稳定。torch.set_grad_enabled(train)在验证时关闭梯度计算显存能省一半左右。loss.item() * imgs.size(0)是按样本数加权因为最后一个 batch 可能不满直接平均会有偏差。参数说明batch_size32是通用起点显存不够降到 16 或 8。lr1e-4是微调的安全值从零训可以用 1e-3。weight_decay1e-4是 L2 正则抑制过拟合。StepLR每 5 个 epoch 把学习率乘 0.5后期收敛更稳。num_workers4是数据加载进程数Windows 上如果报错就设 0。4.2 学习率、batch size 与数据增强的调参逻辑这三个参数是训练效果的主要杠杆但很多人调参靠玄学。说清楚它们各自的作用学习率决定每步更新幅度。太大 loss 震荡不下降太小收敛慢甚至卡在局部最优。微调场景下预训练权重已经很好了学习率必须小1e-4 到 1e-5 之间试。判断方法看前几个 epoch 的 loss如果一直不降降一个数量级如果 loss 变成 NaN说明太大。batch size影响梯度估计的稳定性和显存占用。大 batch 梯度更稳但泛化可能略差小 batch 有正则效果但训练慢。32 是平衡点显存够可以上 64不够降到 16。注意 batch size 变了学习率通常也要跟着调经验是大 batch 配大学习率。数据增强是防过拟合最有效的手段。猫狗数据集里猫狗姿态、光照、背景差异大增强能显著提升泛化。除了翻转和旋转还可以加ColorJitter调亮度对比度、RandomResizedCrop随机裁剪缩放。但增强不是越多越好验证集绝对不能加训练集加太多会让训练变难、收敛变慢。train_tf transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), # 随机裁剪缩放 transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), # 轻度颜色扰动 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])scale(0.7, 1.0)表示裁剪面积占原图的 70% 到 100%太小会丢失目标。ColorJitter的系数别设太大0.2 左右是轻度扰动超过 0.5 可能让图像失真影响学习。5. 训练过程避坑与常见问题排查5.1 验证集准确率虚高的三个来源现象训练时验证集准确率 99%但拿新图片推理一塌糊涂。原因一验证集和训练集有重叠。重组数据时如果没去重同一张图可能同时出现在两边。解决切分前对文件名或图片哈希去重。原因二验证集用了训练时的数据增强。RandomHorizontalFlip这类随机变换如果加在验证集上每次评估输入都不同指标不可比。解决验证集只用 Resize ToTensor Normalize。原因三BatchNorm 在验证时统计量不对。如果验证 batch 太小BN 用 batch 内统计量会失真。解决确保model.eval()被调用验证 batch 别小于 16。5.2 loss 不下降或变 NaN 的排查顺序现象训练几个 epochloss 一直在 0.69 附近二分类的随机水平或者直接变 NaN。排查顺序先看学习率1e-4 还 NaN 就降到 1e-5再看数据标签用ImageFolder时确认class_to_idx和你的预期一致打印几个 batch 的标签看看然后看输入归一化Normalize 的均值方差是否和预训练匹配最后看 loss 函数二分类用CrossEntropyLoss配 2 维输出不要用BCELoss配 2 维那是给 1 维输出用的。5.3 显存不足OOM的四种降级方案现象RuntimeError: CUDA out of memory。按代价从低到高第一降 batch size32 降到 16 再降到 8第二验证时用torch.no_grad()或set_grad_enabled(False)能省不少第三用torch.cuda.empty_cache()清理缓存但治标不治本第四换更小的模型ResNet50 换 ResNet18。如果都不行考虑混合精度训练torch.cuda.amp能把显存占用降一半左右。5.4 训练集准确率高但验证集低过拟合的处理现象训练准确率 99%验证卡在 88% 不动。这是典型过拟合。处理手段按优先级加数据增强最有效、加 weight_decay、加 Dropout、早停验证 loss 连续几个 epoch 不降就停、减小模型。猫狗数据集 2 万多张ResNet18 微调一般不会严重过拟合如果出现了先检查验证集是不是太小或者分布和训练集差异大。5.5 DataLoader 多进程在 Windows 上报错现象num_workers 0时在 Windows 上抛RuntimeError或卡死。原因是 Windows 的进程启动方式和 Linux 不同多进程加载数据需要把训练代码放在if __name__ __main__:保护块里。解决要么加保护块要么把num_workers设成 0。设 0 会慢一些但最省心调试阶段建议先设 0。6. 模型保存、推理与 ONNX 导出把成果用起来训练完只是半成品能推理、能部署才算闭环。先说保存PyTorch 有两种保存方式保存整个模型和只保存参数。推荐只保存 state_dict因为保存整个模型依赖类定义换环境容易加载失败。# 保存最佳模型参数 torch.save(model.state_dict(), best_resnet18_catdog.pth) # 加载时先重建结构再灌参数 model build_model(strategyhead_only) model.load_state_dict(torch.load(best_resnet18_catdog.pth, map_locationcpu)) model.eval()map_locationcpu很重要在有显卡的机器上保存、在没显卡的机器上加载时不加这个参数会报错。推理单张图片的完整流程from PIL import Image def predict(img_path, model, class_names[cat, dog]): img Image.open(img_path).convert(RGB) # 强制三通道防灰度图报错 tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) x tf(img).unsqueeze(0) # 加 batch 维度 with torch.no_grad(): logits model(x) prob torch.softmax(logits, dim1) idx prob.argmax(1).item() return class_names[idx], prob[0][idx].item() label, conf predict(test.jpg, model) print(f预测: {label}, 置信度: {conf:.4f})convert(RGB)是必须的用户上传的图可能是灰度或带 alpha 通道不转会在 ToTensor 后通道数不对。unsqueeze(0)加 batch 维度模型要求输入是[N, C, H, W]。torch.no_grad()关闭梯度推理更快更省显存。如果要把模型部署到非 Python 环境搜「pytorch转onnx」的人就是这个需求。导出代码dummy torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy, catdog.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version11 )dynamic_axes让 batch 维度可变否则导出的模型只能接受固定 batch。opset_version11兼容性好别盲目追新版本。最后说一个我自己的习惯每次实验都固定随机种子把配置写进一个 dict 存下来连同模型一起保存。这样过两周回头看能准确复现当时的实验不用靠回忆。猫狗分类这个项目跑通一遍大概两三个小时但把上面这些坑都踩一遍再填上你对 PyTorch 图像分类全流程的理解会比看十篇教程都扎实。希望帮到你。本文还有配套的精品资源点击获取
返回列表