ARTICLE DETAIL

资讯详情

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

PyTorch实现猫狗公鸡图像分类:CNN实战全流程解析

PyTorch实现猫狗公鸡图像分类:CNN实战全流程解析 简介面向具备一定深度学习基础、希望快速上手PyTorch的初学者这份资源以猫狗公鸡三类图像识别为实战场景完整覆盖数据预处理、CNN网络搭建、训练验证、模型保存加载与结果可视化等关键环节。包内共1390个文件核心是1362张JPG图像数据另有11个Python源码、ONNX导出模型、XML配置、TXT说明与PNG示意图等压缩包整体约554.92MB目录组织清晰便于按步骤查找与复现。目前已有1378人学习下载。借助这份项目包读者不仅能深入理解卷积层、池化层、全连接层在图像分类中的具体作用还能掌握DataLoader数据加载、CrossEntropyLoss损失函数、SGD/Adam优化器、torch.save/torch.load等PyTorch核心API的实操方法从随机翻转、旋转等数据增强到混淆矩阵分析均有完整代码支撑适合在CPU设备上直接完成模型训练与推理为后续深度学习项目打下扎实基础。1. 利用 PyTorch 搭一个猫狗公鸡分类网络一次完整的 CNN 实战复现经常有人问我初学 PyTorch 到底做什么项目才算「真正入门」我的答案是别去跑那些几十层的开源大模型就找一个小而全的三分类任务把数据加载、CNN 搭建、训练验证、模型保存、预测可视化全流程走一遍。这份猫狗公鸡分类项目正好就是干这个的——它用 PyTorch 的nn.Module搭了一个经典 CNN把 1200 张训练图分成猫、狗、公鸡三类并在 CPU 上完成了训练和推理。整个过程不依赖 GPU不碰分布式代码量控制在能看懂、能改、能跑的水平。适合刚学完 PyTorch 基础语法、想动手实践分类任务的开发者也适合那些想快速验证某个网络结构改动效果的工程师拿这个小数据集做基准测试。这个项目的价值不在于网络有多深、精度有多高而在于它把 PyTorch 日常开发里最常用的一套工作流完整地串了起来。你会发现很多在你自己的工作里反复用到的东西——DataLoader的数据组织形式、CrossEntropyLoss的输入格式、模型保存的state_dict方式、以及在 CPU 上推理时的设备切换——全部被压缩在这个小项目里。我今天就把它拆开从数据预处理讲到模型训练再讲几个我在复现时踩过的坑。2. 图像分类的思路与 CNN 选型为什么这个小项目用「三层卷积 全连接」就够了2.1 三分类任务的特征提取逻辑猫、狗、公鸡这三类图像从像素层面看差异很大公鸡有红色鸡冠和鲜艳羽毛猫有胡须和圆耳朵狗有下垂的耳朵和不同的脸部比例。这种差异意味着模型不需要非常深的网络就能学到可区分的特征。和 ImageNet 那种动辄 1000 类、类间差异很小的任务不同三分类任务里浅层 CNN 提取到的边缘、纹理、颜色特征已经足够。这里我建议先理解 CNN 提取特征的层次性。第一层卷积通常只能学到边缘和颜色斑点第二层卷积能组合出纹理和局部形状比如鸡冠的红色区域、猫耳的圆弧轮廓第三层卷积能学到更全局的形态特征。池化层MaxPool2d的作用是逐步缩小特征图尺寸保留主要特征的同时减少计算量并且让模型对轻微的平移有一定鲁棒性。全连接层则把所有特征图展平后的信息映射到三个类别得分上。选择三层卷积而不是更多层有实际考量。1200 张训练图片在数据量上并不大如果网络过深比如直接上 ResNet18 或 ResNet34模型容量远超数据能提供的信息量反而容易过拟合。小数据集配小网络是工程实践中的常用原则。当然如果你的图片数量到了几万张那完全可以考虑迁移学习或更深的网络但这已经不是这个项目讨论的范围了。2.2 数据预处理归一化、尺寸统一和数据增强参数项目里的训练数据是 1200 张图片验证集数据量是 300 张左右。在把图片喂给网络之前必须先做三件事统一尺寸、数值归一化、划分数据集。尺寸统一很好理解因为 CNN 的全连接层要求输入维度固定。项目中设置的是 224×224这是常见的选择不过也可以改成 128×128 来加速训练。归一化方面常见做法是把像素值从 [0,255] 缩放到 [0,1]再按 ImageNet 的平均值和标准差做标准化。这样做的原因是如果像素值范围过大损失函数的梯度更新会不太稳定而且 PyTorch 预训练模型如果存在的话一般也是按标准化后的输入设计的。数据增强是提升泛化能力的重点。常见的做法包括随机水平翻转RandomHorizontalFlip、随机旋转RandomRotation和随机裁剪RandomResizedCrop。对于公鸡这类有明显朝向的物体水平翻转不会破坏语义翻转后的公鸡还是公鸡所以可以放心用。我一般会加一个 10 度以内的随机旋转稍微增加一点多样性就够了旋转角度过大会导致物体变形失真。2.3 划分训练集和验证集我习惯用 8:2 的比例划分训练集和验证集。1200 张图的场景一般不需要单独留测试集因为验证集已经足够反映模型在未见数据上的表现。实现时可以用torch.utils.data.random_split这个函数接受一个数据集实例和一个比例列表from torch.utils.data import Dataset, DataLoader, random_split # dataset 是你的自定义 Dataset 实例 train_size int(0.8 * len(dataset)) val_size len(dataset) - train_size train_dataset, val_dataset random_split(dataset, [train_size, val_size])random_split的好处是不用手动索引数组而且它默认是随机切分的避免了因为图片按类别顺序存放而导致的分布偏差。需要注意的一点是拆分后train_dataset和val_dataset不共享同一份数据训练时数据增强只应该加在训练集上。验证集应该保持原始输入这样才能客观评估模型的真实性能。我见过有人在验证集上也做了随机翻转结果验证精度忽高忽低最后发现是数据增强干扰了评估——这就是典型的坑。2.4 自定义 Dataset 的写法这个项目的数据组织方式是典型的从文件夹读取图片每个类别的图片放在各自的子目录里。这时候直接用torchvision.datasets.ImageFolder是最省事的from torchvision import datasets, transforms transform_train transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # ImageFolder 会自动按文件夹名字排序依次分配 class index train_set datasets.ImageFolder(rootdata/train, transformtransform_train) val_set datasets.ImageFolder(rootdata/val, transformtransform_val)这里ImageFolder会自动读取每个子目录下的图片并且按目录名排序分配类别索引。比如目录名是cat、dog、cock那么索引就是0: cat、1: dog、2: cock。这个顺序要注意因为后面预测时候需要知道索引对应的真实类别名往往是新手容易搞混的点。如果你的数据不是按文件夹组织的而是像项目里那样散落在根目录那就需要自己写一个Dataset子类核心是重写__getitem__方法和__len__方法。文件列表读取用glob.glob或者自己维护一个 CSV 文件都行。3. 从零搭建 CNN 模型nn.Module 的结构设计与各层参数详解3.1 网络结构卷积层、池化层、全连接层的组合方式有了预处理好的数据下一步就是定义模型。这个项目的核心就是nn.Module它是 PyTorch 中所有神经网络模块的基类。我们的 CNN 采用经典结构特征提取部分由三个卷积块堆叠每个卷积块包含Conv2d、ReLU、MaxPool2d最后接展平操作和三个全连接层输出三类的得分。import torch.nn as nn import torch.nn.functional as F class CnnClassifier(nn.Module): def __init__(self, num_classes3): super(CnnClassifier, self).__init__() # 第一个卷积块输入 3 通道RGB输出 16 个特征图 self.conv1 nn.Conv2d(3, 16, kernel_size3, padding1) self.bn1 nn.BatchNorm2d(16) self.pool1 nn.MaxPool2d(2) # 第二个卷积块16 - 32 个特征图 self.conv2 nn.Conv2d(16, 32, kernel_size3, padding1) self.bn2 nn.BatchNorm2d(32) self.pool2 nn.MaxPool2d(2) # 第三个卷积块32 - 64 个特征图 self.conv3 nn.Conv2d(32, 64, kernel_size3, padding1) self.bn3 nn.BatchNorm2d(64) self.pool3 nn.MaxPool2d(2) # 分类部分展平后接全连接层 self.fc1 nn.Linear(64 * 28 * 28, 256) self.fc2 nn.Linear(256, num_classes) def forward(self, x): x self.pool1(F.relu(self.bn1(self.conv1(x)))) x self.pool2(F.relu(self.bn2(self.conv2(x)))) x self.pool3(F.relu(self.bn3(self.conv3(x)))) x x.view(x.size(0), -1) # 展平保留 batch 维度 x F.relu(self.fc1(x)) x self.fc2(x) return x代码里有几个细节需要展开说。padding1是为了让卷积操作不改变特征图的宽高。以 224×224 输入为例经过conv1后尺寸保持 224×224再经过pool1变成 112×112。连续三轮之后到全连接层之前特征图大小是 28×28所以全连接层的输入维度是64 * 28 * 28。如果你改了输入尺寸这一项就要跟着重新算很多人在这一步翻车都是因为算错了维度。BatchNorm2d的加入是因为这个小数据集上训练时网络容易受每批数据分布差异的影响。BatchNorm 将每个 batch 的激活值归一化到均值为 0、方差为 1然后用可学习的缩放因子和偏移量进行变换。它的效果是加速收敛同时让网络对学习率的敏感度降低。项目原文没有提到 BatchNorm但我强烈建议加上因为这是用最少的代码换取稳定训练收益的手段。不加的话模型可能需要调整到更低的学习率才能稳定训练。3.2 forward 背后的自动求导机制在forward函数之外PyTorch 的自动求导是理解 CNN 训练的关键。nn.Module内部注册了所有的参数训练时调用loss.backward()会自动计算梯度不需要手写反向传播。你只需要确保forward中所有操作都是可微的——Conv2d、ReLU、MaxPool2d、BatchNorm、Linear 都是标准的可微操作所以整个网络没问题。这里有一个容易被忽略的点forward函数里执行x.view(x.size(0), -1)时必须保留 batch 维度。如果写成x.view(-1)会把整个 batch 的数据拉成一个一维向量后续的Linear层维度完全对不上运行时报错是肯定的。这件事看起来基础但实际写代码时非常容易写错特别是在网络结构调试阶段报错信息里出现mat1 and mat2 shapes cannot be multiplied时多半就是这个view的问题优先检查全连接层的输入维度是否和展平后的特征图维度匹配。3.3 初始化与设备切换模型搭建完成后训练之前要做两件事初始化权重、移动到计算设备上。PyTorch 会默认用 Kaiming 均匀分布初始化卷积层权重一般情况下可以直接用默认值。但如果你发现模型训练初期 loss 不掉可以尝试手动初始化def init_weights(m): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) elif isinstance(m, nn.Linear): nn.init.xavier_normal_(m.weight) nn.init.zeros_(m.bias) model CnnClassifier(num_classes3) model.apply(init_weights)设备切换方面需要注意 CPU 和 GPU 环境的兼容问题。项目明确说明最终在 CPU 上做预测因此模型保存和加载时要保持一致。如果训练时用了 GPU保存模型后想放到 CPU 上推理就必须在加载时指定映射。# GPU 上训练CPU 上推理的标准做法 model CnnClassifier(num_classes3) model.load_state_dict(torch.load(model.pth, map_locationtorch.device(cpu))) model.eval()map_locationtorch.device(cpu)这个参数特别关键。PyTorch 保存的权重如果来自 CUDA 张量直接加载到 CPU 会报设备不匹配错误。很多人在这出错后完全摸不着头脑其实一行代码就解决了。另一件容易忘的事是加载模型后忘记调model.eval()。eval()会关闭 BatchNorm 的批次统计更新和 Dropout 的随机丢弃如果忘记调用同一个输入跑两次得到的结果可能不一样这就是典型的「训练时好好的推理时结果不对」的坑。4. 训练循环与验证评估从损失函数到模型保存的完整闭环4.1 损失函数和优化器的选择三分类任务用CrossEntropyLoss这是完全正确的选择。PyTorch 的CrossEntropyLoss内部已经包含了LogSoftmax和负对数似然损失所以模型的输出不需要额外套一层Softmax。模型输出的三个原始得分logits直接喂给损失函数即可。这一点我见过不少初学者弄错在forward的返回值里加了F.softmax然后在损失函数里又过了一遍CrossEntropyLoss结果梯度计算路径被破坏训练完全跑偏。优化器方面SGD 和 Adam 都行。在这个小数据集上我实测更推荐用 Adam 或 AdamW因为它的自适应学习率机制让超参数调优成本低很多。用 SGD 需要手动调 momentum、weight decay、学习率衰减策略项目目标是快速实现分类没必要把时间花在这上面。默认lr0.001就够用了。import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001, weight_decay1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1)weight_decay1e-4是 L2 正则化这在 1200 张图的小数据集上很重要可以有效防止过拟合。StepLR表示每 10 个 epoch 把学习率乘以 0.1。这个学习率衰减策略不能省因为 Adam 在训练后期如果一直保持较大学习率容易在最优值附近震荡导致 loss 曲线不收敛。4.2 训练循环的代码模板数据加载这一环DataLoader的参数值得多聊两句。batch_size我建议设置 16 或 32。CPU 训练时batch_size太大内存占用和单次迭代时延都会飙高太小则梯度更新不稳定。num_workers在 Windows 上建议设 0在 Linux 上可以设 2 到 4 来提高数据读取吞吐。from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers2)shuffleTrue必须在训练集上开启否则每个 epoch 看到的数据顺序都一致模型会学到数据顺序的伪规律。验证集不需要 shuffle而且num_workers和训练集保持一致即可。写完数据加载下面就是核心的训练循环。def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() total_loss 0 correct 0 total 0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) return total_loss / total, correct / total训练循环的标准流程是清零梯度 → 前向传播 → 计算损失 → 反向传播 → 更新权重。需要特别强调的是optimizer.zero_grad()不能被忽略。PyTorch 默认是累积梯度如果不清零每个 batch 的梯度会累加到上一步等效于用一个极大的 batch_size 训练几乎必然导致 loss 异常升高或直接爆炸。至于对inputs和labels调用.to(device)是因为如果模型在 GPU 上而数据在 CPU 上运行时会报错设备不匹配。即使你只打算用 CPU也建议保留这行代码以便将来换了带 GPU 的环境可以直接训练。4.3 验证与早停策略验证流程比训练简单关键区别是要在torch.no_grad()上下文里跑这样可以省掉梯度计算的内存开销同时也不会因为验证集上的反向传播而污染模型状态。另外model.eval()必须放在验证循环前。def evaluate(model, val_loader, criterion, device): model.eval() total_loss 0 correct 0 total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) total_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) correct (predicted labels).sum().item() total labels.size(0) return total_loss / total, correct / total训练时每 epoch 跑一次验证同时记录验证 loss当验证 loss 连续若干 epoch 不再下降时就把当前模型权重保存为最佳模型。不要只看验证准确率因为准确率在某些情况下会出现波动而 loss 的下降趋势更为平滑、稳定。我在项目里通常记录两个指标保存的最佳模型以验证 loss 为准。这里有一点经验训练集准确率接近 100% 但验证准确率只有六七十的时候说明模型过拟合了此时优先做数据增强、降低模型复杂度或增加 weight_decay而不是盲目继续训练。4.4 模型保存和加载的推荐写法PyTorch 保存模型有几种方式训练全部参数、结构、优化器状态的完整快照可以用torch.save(model, model.pth)但我不推荐这种方式因为加载时要求环境和保存时完全一致Python 类名、模块路径变了就会报错。推荐的做法是只保存state_dict# 保存 torch.save(model.state_dict(), model.pth) # 加载 model CnnClassifier(num_classes3) model.load_state_dict(torch.load(model.pth))使用state_dict方式保存的好处是加载时只需要重新实例化模型结构然后填入权重即可。这种方式对迁移学习和模型共享都更友好。如果除了模型权重还想保存一些训练状态信息比如优化器状态、epoch 数、验证精度可以打包成一个字典checkpoint { model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, val_acc: best_val_acc } torch.save(checkpoint, checkpoint.pth)我的习惯是每个 epoch 结束后都保存一次 checkpoint而不是只有训练结束时才保存。这样如果训练中途因为意外中断可以随时从最近的 checkpoint 恢复相当于一份「后悔药」。恢复时只需要加载字典里的各部分再手动把模型调到训练模式即可。5. 避坑与排查CPU 训练环境下最常遇到的五个问题5.1 数据加载慢训练一个 epoch 要半小时现象CPU 训练时每个 epoch 的大部分时间都花在数据读取上GPU 利用率很低在 CPU 场景下表现为 CPU 时间全花在 IO 上训练速度极慢。原因图片解码和尺寸调整是纯 CPU 操作。torchvision的默认图像后端 PIL 是单线程解码大量图片连续读取时会成为瓶颈。num_workers0时数据只能串行加载速度受限于单核性能。解决Linux 上设置num_workers4或更高让子进程并行加载数据。Windows 上的多进程数据加载有兼容性问题建议先用较小num_workers测试不行就设 0。另外如果图片都是很大的原图可以先离线把所有图片统一尺寸并另存为小分辨率的副本再参与训练。这样能显著减少训练时的图像解码开销。5.2 模型输出 NaNloss 变成无穷大现象训练进行到某个 epoch 后loss 突然变成nan而且后续无法恢复。原因最常见的原因是学习率过大导致梯度爆炸权重变成nan后无法恢复。其次是数据集中有损坏的图片或者像素值包含异常值。还有一种情况是 BatchNorm 在 batch_size 太小时方差为 0产生除零错误。解决先把lr调到0.0001验证是否能稳定训练再用torch.isnan检查输入数据。同时尽量把batch_size保持在 16 以上避免 BatchNorm 统计量不稳定。排查时建议在损失函数附近加一个断言一旦出现nan就立刻停止训练并打印当前层权重。5.3 训练集精度高、验证集精度低模型严重过拟合现象训练集准确率到 95% 以上但验证集只有 60% 左右且验证集的 loss 随训练进度不降反升。原因模型容量过大、训练样本不足、缺少数据增强或正则化导致的典型过拟合。1200 张图片对三层 CNN 来说虽然不算剧烈但如果不加数据增强特别是随机翻转和随机裁剪模型很快就会记住训练图的具体像素分布。解决在训练集上开启RandomHorizontalFlip和RandomRotation(10)同时把weight_decay从1e-4调大到1e-3或者减少特征图数量比如把三层的通道数从 64 减到 32缩小模型容量。过拟合是「欠拟合平衡」问题需要同时调整数据和模型两侧。5.4 加载模型预测时报错size mismatch现象模型加载state_dict时抛出size mismatch for fc1: copying a param with shape ...这样的错误。原因你实例化的模型结构和保存模型时的结构不一致。常见情况是保存时用的num_classes3加载时写成了num_classes10全连接层的输出维度对不上。解决加载前务必核对模型类的定义和初始化参数。建议将num_classes这样的配置提取成常量在训练脚本和推理脚本中共用避免脚本之间手动同步数字带来的隐患。另外保存模型文件的命名建议带参数量信息比如cnn_catdogcock_3cls.pth减少误用。5.5 验证时忘记model.eval()结果时好时坏现象同一个输入图片连续预测多次结果会变来变去有时输出不同类别。原因模型处于训练模式。BatchNorm 层会在每次前向传播时用当前 batch 的统计数据更新全局统计量导致权重不变但输出却因为统计量变化而改变。如果网络里有 Dropout 层问题更严重它会随机屏蔽神经元直接改变输出结果。解决在推理之前所有模型的参数会影响前向传播的层请统一调用model.eval()。这是最容易修复但有最容易漏掉的问题。我习惯在写推理函数时第一行就写上model.eval()和model.load_state_dict绑定在一起。6. 进阶混淆矩阵、单图验证与 GPU 迁移的改进方向项目做到这里基础分类已经能跑通了。但要把这个小项目变成可交付的工程还需要补两块内容一是结果的可视化和分析二是从 CPU 到 GPU 的迁移准备。6.1 绘制混淆矩阵找到模型到底在哪些类别上犯错准确率只有一个数字无法告诉你模型到底把狗认成了猫还是把公鸡认成了猫。用一个sklearn的confusion_matrix就够了这是我在分类项目里每次都会做的基本分析。import matplotlib.pyplot as plt import numpy as np from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay model.eval() all_preds [] all_labels [] with torch.no_grad(): for inputs, labels in val_loader: inputs inputs.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) disp ConfusionMatrixDisplay(confusion_matrixcm, display_labels[cat, dog, cock]) disp.plot(cmapBlues) plt.savefig(confusion_matrix.png, dpi150)拿到混淆矩阵后怎么读重点看对角线之外的数字。如果「狗→猫」的错误明显多于其他错误说明狗的某些纹理特征和猫的过于接近可以针对性做数据增强比如增加这两类图片的亮度扰动如果是「公鸡→猫」的错误较多大概率是背景噪声干扰可能需要考虑对图片做裁剪来突出主体。混淆矩阵能把你从「不知道模型错在哪」的黑匣子状态里拉出来。从实验结果来看这个三分类模型大概率在公鸡和狗之间的混淆较少因为它们外观差异大最容易混淆的通常是猫和狗的耳朵形状、毛色深浅等局部特征。如果出现这类情况不需要改网络结构增加数据增强的强度或者增加训练 epoch 往往就能改善。6.2 单张图片的独立预测脚本训练验证结束后写一个独立的单图预测脚本方便对项目附件里那些散图直接做测试。这里有一个细节预测时输入图片必须和应用在训练集上的变换完全一致尺寸、归一化参数都要对齐否则预测结果没有意义。from PIL import Image def predict_image(model, image_path, device): model.eval() img Image.open(image_path).convert(RGB) 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]) ]) img_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(img_tensor) _, pred torch.max(outputs, 1) class_names [cat, dog, cock] return class_names[pred.item()]注意unsqueeze(0)是给图片加一个 batch 维度因为模型接受的是四维张量[batch_size, channels, height, width]。忘记这一步会直接报维度错误。另外预测时不需要Softmax也可以直接torch.max拿到最大概率的类别索引因为argmax的结果在Softmax前后是等价的这是一个常被误用的点。6.3 从 CPU 迁移到 GPU 的几个改动点如果你之后想换到 GPU 上训练改动非常小。除了把模型和数据都.to(cuda)还要注意几个细节。第一torch.load加载模型到 CPU 时要用map_locationtorch.device(cpu)否则 PyTorch 默认回到模型权重保存时的设备。第二DataLoader的num_workers在 GPU 环境下可以适当调高到 4 或 8因为 GPU 训练速度更快时CPU 的数据加载必须跟上。第三训练过程中定时torch.cuda.empty_cache()防止显存碎片累积。CPU 环境下的训练速度确实比 GPU 慢不少但这个数据集规模小、网络层数浅CPU 也可以在合理时间内完成训练。以 1200 张图片、batch_size 32、20 个 epoch 为例普通 CPU 大概在十几分钟到半个小时内完成一轮完整训练完全在可接受的范围内。如果你手里的机器性能一般可以先把输入尺寸改成 128×128训练时长会成倍缩短精度损失在这个分类任务下不会太明显。说回项目本身PyTorch 这套工作流的掌握关键从来不是某个单独的 API而是把「数据 → 模型 → 训练 → 评估 → 部署」这条路走通。从这以后我每次搭分类模型都会强制走一遍完整闭环先写数据检查脚本确认标注和图片对得上再跑一个过拟合小样本来验证网络能学最后才上完整训练和混淆矩阵分析。做完这个猫狗公鸡项目你会发现自己很快就能把同样的套路迁移到其他分类场景里——换个数据集改一下num_classes基本就完事了。希望这篇拆解帮到你。本文还有配套的精品资源点击获取
返回列表