ARTICLE DETAIL

资讯详情

深耕商务建站与企业官网运营的一线实战洞察。

PyTorch实战:从零构建食物图片分类器的完整记录

PyTorch实战:从零构建食物图片分类器的完整记录 从零搭一个食物图片分类器我的 PyTorch 入门实战记录想入门 PyTorch但又不想只跑个 MNIST 就完事那食物图片分类是个特别合适的练手任务。它既不像 MNIST 那样简单到没有成就感也不像 COCO 检测那样一上来就需要分布式训练和多卡策略难度刚好卡在一个“跳一跳够得着”的位置。这篇文章是我自己从零开始实现一个食物图片分类项目的完整记录包含数据集整理、预处理细节、模型选型、训练循环手写、调参经验、踩坑复盘。项目本身不复杂但能把 PyTorch 的日常工作流全部过一遍。做完之后你会发现后面再去碰迁移学习、目标检测、语义分割这些任务底子已经打得差不多了。适合谁看刚学完 Python 基础、看过几篇 PyTorch 教程但没完整做过项目的朋友以及想用一份清晰代码把“训练-评估-预测-模型保存”整条链路跑通的人。如果你是资深炼丹师这篇文章可能偏基础但里面有几个关于数据加载和调参的坑也许能让会心一笑。1. 项目整体拆解食物分类到底在解决什么问题1.1 食物分类任务的核心难点食物图片分类本质上是一个图像分类问题模型要做的就是输入一张图片输出它属于哪个类别。我这边定了 10 个常见食物类别面包、乳酪、冰淇淋、苹果、香蕉、玉米、胡萝卜、煎蛋、米饭、意面。听起来和 CIFAR-10 很像但实际做起来有几个地方和普通分类不一样。第一是类内差异大。同样是米饭有炒饭、白米饭、盖浇饭颜色、纹理、摆盘方式完全不同。同一个类别里的图片可能比不同类别之间的差异还大这对模型的特征提取能力提出了要求。第二是类间相似度高。胡萝卜和玉米都是长条形的苹果和番茄都是圆形的仅仅靠颜色或者形状一个维度的特征根本分不开。这也是为什么传统手工特征做食物分类效果一直不好的原因需要模型自己去学更抽象的组合特征。第三是数据噪声多。食物图片来自真实生活场景有餐盘、桌面、餐具、手指、光影干扰甚至有咬过一口的苹果。如果只用干净的商品图训练到了真实场景基本就崩了所以数据增强在这个任务里不是可选项而是必需品。1.2 为什么选择 PyTorch 而不是其他框架做这个项目我选 PyTorch不是因为别的框架不行而是因为 PyTorch 的调试体验对新手最友好。最核心的一点是动态计算图。模型跑起来之后你可以随时在任意一行代码处打断点查看中间张量的 shape、dtype、梯度值改完代码不用重新构建整个图就能继续跑。这种交互式的开发体验对新手理解“前向传播-反向传播”的整个过程特别有帮助。你自己手写一个训练循环比用任何高层封装都更能理解框架内部发生了什么。另外 PyTorch 的生态在视觉领域是最完整的。torchvision 模型库、transform 工具、预训练权重、各类第三方扩展基本上你能想到的视觉处理需求都有现成的轮子。社区讨论也多遇到问题搜一下大概率能找到答案。1.3 整体技术选型与管线设计整个项目按下面这条管线走数据采集与整理 - 预处理 Pipeline - DataLoader 构建 - 模型选型 - 训练循环实现 - 评估与保存 - 单张图片推理硬件方面我用的是一块显存 8GB 的中端显卡。这个配置训练 10 分类的 ResNet18 完全够用Batch Size 开到 64 也没压力。没有 GPU 也能跑就是慢一些CPU 模式下建议用更小的模型或者把图片尺寸降下来。软件环境方面PyTorch 版本用的 2.1Python 3.10torchvision 0.16。框架内部的接口这两年有些变化我代码里会加一点兼容性说明如果你用的版本不一样大概率也能对上。2. 数据集准备与预处理这步做不好后面全白搭2.1 数据集来源与整理方式数据是分类任务的命根子。我没有直接用某个现成的食物数据集而是从一个开源图片平台上抓取标注好的食物图片再结合一部分自己手拍的图片做补充。整理成下面这样的目录结构food_dataset/ ├── train/ │ ├── bread/ │ ├── cheese/ │ ├── ice_cream/ │ ├── apple/ │ ├── banana/ │ ├── corn/ │ ├── carrot/ │ ├── fried_egg/ │ ├── rice/ │ └── spaghetti/ └── val/ ├── bread/ └── ...每个类别我收集了大约 300 张训练图和 80 张验证图。这个数据量对于初学者完全够用了如果再少可能需要更强的数据增强或迁移学习来兜底。目录按类别分好之后用 torchvision.datasets.ImageFolder 就可以直接读取它会把子目录名自动映射成类别标签。这里有个关键点子目录的命名尽量用英文避免编码问题导致读取错误因为某些系统对中文路径默认编码是 GBKPyTorch 底层用的是 UTF-8碰到中文目录名容易临时找不到文件。2.2 预处理 Pipeline 的几个关键细节预处理是整个项目最容易忽略但影响最大的部分。torchvision 的 transform 流水线我这样定义from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.6, 1.0)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), 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_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]) ])训练集用了 RandomResizedCrop这一步很关键它不仅仅是裁剪还会随机改变图片的缩放比例和宽高比相当于做了模拟距离变化和拍摄角度的增强。scale 参数我设成 (0.6, 1.0)意思是裁剪区域占原图的 60% 到 100%。这个区间对食物比较合适太低了把整个食物裁掉一半反而引入无效噪声。归一化的 mean 和 std 直接用 ImageNet 的统计值因为我们后面会加载 ImageNet 预训练权重输入分布必须和预训练时一致否则前期微调阶段梯度会不稳定收敛慢很多。验证集一定不要用 RandomResizedCrop。很多新手在验证集也上随机裁剪导致每次评估同一张图都可能得到不同的预测结果最后模型得分忽高忽低根本没法判断好坏。验证集用 Resize 到 256 再 CenterCrop 到 224这个流程是 ImageNet 分类任务的标配好处是既保留中心主体信息又统一了输入尺寸。2.3 数据划分与 DataLoader 的优化配置数据划分我是直接在目录结构层面完成的训练和验证独立文件夹简单可靠。没有做 K 折交叉验证因为对初学者来说单次划分配合早停法已经足够判断模型好坏。DataLoader 的配置里有几个参数值得注意from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue)num_workers 设为 4在我的机器上数据加载和训练基本可以并行GPU 不会因为等数据而空转。这个参数不是越大越好太多进程反而会增加 CPU 调度开销和内存占用。经验值一般是 CPU 核心数的一半左右如果你的机器是 8 核设 4 是安全的起点。pin_memoryTrue 也很重要它把数据加载进锁页内存从 CPU 拷贝到 GPU 时会走更快的 DMA 通道。代价是占用一部分固定内存但性能收益在数据量大的时候非常明显。这里有个小坑如果你的内存只有 8GB数据又特别大的情况下pin_memory 可能导致内存压力过大那就别开。3. 模型构建从手写网络到 ResNet 迁移学习3.1 模型结构选型与对比模型选型我认真对比过三个方案。第一个方案是自己写一个三层卷积网络结构像这样import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128 * 28 * 28, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, x): return self.classifier(self.features(x))这个网络在 CIFAR-10 这种小图上可以跑一跑但换到 224x224 的食物图片上三层卷积的视野范围太小提取不到足够抽象的语义特征训练两轮就能看到验证准确率卡在 60% 左右上不去。这就是典型的模型容量不足。第二个方案是用 torchvision 里现成的 ResNet18from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, 10)这里最关键的一行是 model.fc nn.Linear(num_ftrs, 10)它把 ResNet18 原来的 1000 类分类头换成了我们自己的 10 类分类头。前面的卷积层依然保留原来的参数这些参数已经在 ImageNet 上学会了通用的边缘、纹理、形状特征相当于省掉了从头学特征的时间。第三个方案是 ResNet34。比 ResNet18 多了一层残差块参数量大一些但在这个 10 分类小数据集上提升不明显反而训练速度更慢。最终我选择 ResNet18这个选择非常稳妥预训练权重齐全、显存占用小、推理速度快后续还可以拿它做实时 demo。3.2 迁移学习的三层境界用预训练模型做迁移学习不同阶段策略不一样。最简单的做法是冻结卷积层只训练分类头。具体操作是将前面所有参数的 requires_grad 设为 False只有最后一层保持 Truefor param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True这样做的好处是训练速度极快因为反向传播只经过最后一层梯度计算量小。缺点是特征提取网络是固定的无法针对食物数据做自适应优化。如果数据量少这个方法很管用我在 300 张每类的数据集上试过验证准确率大概能到 75%。更常用的做法是解冻部分网络做全参数微调。做法就是什么都不用改全网络参数都保持可训练但把学习率调小。因为如果从头到尾都用默认学习率 0.01 训练预训练模型前面卷积层的预训练特征会被破坏。我用的是不同层不同学习率的方法fc_params model.fc.parameters() conv_params [param for name, param in model.named_parameters() if name not in {fc.weight, fc.bias}] optimizer torch.optim.SGD([ {params: conv_params, lr: 1e-4}, {params: fc_params, lr: 1e-3} ], momentum0.9)这样卷积层用一个较小的学习率微调分类头用一个较大的学习率从头学。经验证明这样做比全局统一学习率效果更好能兼顾特征适配和分类器收敛速度。还有一个进阶做法叫层级解冻就是把网络从后往前逐步放开训练每训练几个 epoch 解冻一层。实际操作中感觉收益性价比不高在数据集不太大的情况下全参数微调 小学习率就已经很好了。3.3 损失函数、优化器与评估指标的选择分类任务最常用的损失函数是交叉熵损失criterion nn.CrossEntropyLoss()这个函数把模型的原始输出 logits 和标签作为输入内部自动完成 softmax 计算和交叉熵运算。需要注意模型输出的是 logits不要在外面再包一层 softmax。新手经常犯这个错误导致训练初期损失值异常地高因为 softmax 之后再算交叉熵数值稳定性会变差。如果需要在推理时得到概率输出那是在模型输出之后单独调 softmax而不是加在模型结构里。优化器我最终选了 SGD 而不是 Adam。很多人觉得 Adam 好收敛快但是 SGD 在配合动量之后泛化能力往往比 Adam 好尤其在数据量不大的情况下。学习率我设成 0.001 起步配合余弦退火调度器。评估指标不能只看准确率。食物分类的类别不均衡问题不严重但为了更稳妥我还是同时监控每个类别的准确率和 Macro F1。准确率容易给人一种“模型已经很好”的错觉因为如果某几个类别特别简单模型会靠它们拉高整体准确率但真正难的类别可能完全没学会。4. 训练循环手写与参数调优模型的成长过程4.1 一个标准训练循环的完整实现我坚持手写训练循环不用任何高层封装。原因很简单训练循环是 PyTorch 的心脏不手写一遍就永远隔着一层纱。完整代码如下import torch device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) best_val_acc 0.0 epochs 30 for epoch in range(epochs): model.train() running_loss 0.0 correct 0 total 0 for images, labels in 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() _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() train_loss running_loss / len(train_loader) train_acc 100.0 * correct / total # 验证阶段 model.eval() val_loss 0.0 val_correct 0 val_total 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() _, predicted torch.max(outputs.data, 1) val_total labels.size(0) val_correct (predicted labels).sum().item() val_acc 100.0 * val_correct / val_total print(fEpoch {epoch1:02d} | Train Loss: {train_loss:.4f} | fTrain Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth)这一段代码里有几个细节必须说明。model.train() 和 model.eval() 很重要。这两个模式切换直接影响 BatchNorm 和 Dropout 的行为。train 模式下 BatchNorm 会使用当前 batch 的统计量Dropout 会随机丢弃神经元eval 模式下 BatchNorm 使用累计的全局统计量Dropout 被关闭。如果不切换到 eval 模式验证时结果会不稳定这是个常见坑。with torch.no_grad() 让验证阶段完全不计算梯度。这能省下大量显存和计算时间。在验证阶段我们只需要前向传播不需要通过 loss.backward() 更新梯度所以直接开启无梯度模式。optimizer.zero_grad() 放在前向之前每次迭代都清空梯度。如果不调用这个梯度会累加在某一步更新的时候用的就是几个 batch 的累积梯度训练过程会变得极度不稳定。PyTorch 不会自动清梯度这一点和有些高层框架不一样必须手动清。验证准确率超过历史最好就保存模型。我每次 epoch 结束都检查一次保存的时候只保留 val_acc 最高的那一次。这种“早停 快照”策略比无条件跑完所有 epoch 更稳妥后面加载到的最佳模型不需要再人工挑选。4.2 超参数调优的完整记录学习率是最关键的超参数。我做了三次对比实验学习率优化器最终验证准确率训练稳定性0.1SGD62%波动大后期发散0.01SGD78%前期震荡后期稳定0.001SGD85%平稳上升稳定收敛0.001Adam84%收敛最快但后期略波动0.1 这个学习率对于迁移学习来说太激进了。因为预训练权重的特征已经很好学习率太大一步跨过头就把原有知识破坏了。0.001 是迁移学习微调阶段比较安全的选择。Batch Size 的影响我也对比了。Batch Size显存占用训练速度稳定性16约 1.5GB最慢震荡剧烈64约 3GB快平稳128约 5GB最快有些微震荡太小 batch 会产生比较大的梯度噪声训练震荡明显。最后我选 64因为它在我的 8GB 显存卡上留下了充足的余量还可以顺便开个并行的验证程序。优化器动量我设为 0.9这是一个经验上几乎不需要调的默认值SGD 加动量能在保证梯度方向稳定的同时加快收敛速度。4.3 过拟合的判断与应对方法训练过程中我在第 5 个 epoch 开始观察到训练准确率接近 90%但验证准确率只有 72%。这就是过拟合的典型信号。食物图片分类数据集量不算大模型容量又大过拟合几乎是必然的。我从两方面来应对。第一是数据增强。项目已经挂了 RandomResizedCrop、RandomHorizontalFlip、RandomRotation、ColorJitter这些都属于基础操作。为了再强化我给训练集加了一个 MixUp 增强import numpy as np def mixup_data(x, y, alpha0.2): if alpha 0: lam np.random.beta(alpha, alpha) else: lam 1 batch_size x.size()[0] index torch.randperm(batch_size).to(device) mixed_x lam * x (1 - lam) * x[index] y_a, y_b y, y[index] return mixed_x, y_a, y_b, lamMixUp 的思路是把两张图片按比例混合标签也跟着混合。比如一个苹果和一张香蕉按 0.7:0.3 混合模型要学到的是这个混合图片有 70% 的苹果和 30% 的香蕉。这样模型被迫学习软标签决策边界更加平滑不容易陷入尖锐的过拟合状态。第二是Dropout 和权值衰减。ResNet18 本身已经带了全局平均池化对于防止过拟合有一定帮助。我在分类头前面加了 Dropout(0.2)并且设置了 weight_decay1e-4也就是 L2 正则化限制权重的绝对值大小。加了这两层保护后验证准确率有 3-4 个百分点的提升。5. 训练过程中踩过的坑与排查实录5.1 显存不足的排查过程第一次跑 ResNet18 加 batch 128大概在训练到第 3 个 epoch 时爆了显存。报错信息是 CUDA out of memory。我的排查步骤是这样的先看是不是常量缓存占用的显存用 torch.cuda.empty_cache() 清一轮再试确认不是缓存问题后把 batch size 降到 64显存占用从 5.2GB 降到 3GB。问题解决。这里有一个额外的技巧如果只是显存峰值爆掉可以开启 PyTorch 2.0 的自动混合精度训练float16 半精度计算会把显存占用差不多减半。代码如下from torch.amp import autocast, GradScaler scaler GradScaler() for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with autocast(device_typecuda): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度对我来说不仅是显存优化更是训练加速手段。我那台显卡的 FP16 计算速度接近 FP32 的两倍整个训练时间缩短了接近一半。初学者不用害怕这个 API 复杂直接用就行内部机制暂时不需要完全搞懂。5.2 训练不收敛的问题排查有段时间模型 loss 一直停留在 2.3 左右准确率始终在 35% 上下像是卡住不动。我排查过后发现问题出在学习率预热缺失。迁移学习微调时如果一开始就用中等偏大的学习率预训练特征会在早期被大幅扰动导致 loss 飙升后无法恢复。解决办法是在前几个 epoch 用很小的学习率做预热from torch.optim import lr_scheduler def warmup_lr(epoch, warmup_epochs3, base_lr0.001): if epoch warmup_epochs: return base_lr * (epoch 1) / warmup_epochs return base_lr scheduler lr_scheduler.LambdaLR(optimizer, lr_lambdawarmup_lr)前 3 个 epoch 学习率从 0 逐渐升到 0.001之后再正常训练。加了这一层预热后模型在第一个 epoch 的 loss 就降到了 1.5 附近训练明显走上了正轨。另一个不收敛的常见原因是标签问题。我检查了一遍标签映射发现 ImageFolder 是按文件名排序的目录名顺序和类别数组顺序可能不一致。排查方法是打印前几个样本的标签和类别名对齐关系确认无误后再训练。5.3 数据加载慢与 CPU 瓶颈训练过程中 GPU 利用率经常掉到 60% 左右看代码逻辑应该是数据加载跟不上。我用 num_workers4 解决了这个问题。注意几个细节workers 不是越多越好。我试过 num_workers8GPU 利用率没明显提升反而程序启动变慢偶尔还会出内存错误。因为我那台机器还要跑其他服务8 个进程把 CPU 占满了。还可以开启验证集数据的 preload把整个验证集一次性加载进内存# 缓存验证集到显存减少重复 IO val_images torch.cat([batch[0] for batch in val_loader], dim0).to(device) val_labels torch.cat([batch[1] for batch in val_loader], dim0).to(device)验证集图片数量少几百张图占不了多少显存但节省每次验证时反复从磁盘读文件的 IO 时间。实际体验是每个 epoch 的验证时间从 15 秒降到了 3 秒。6. 推理部署与结果分析模型到底学到了什么6.1 单张图片推理的完整流程训练完以后我把模型保存成文件然后做了一个简单的推理脚本。推理过程是独立于训练的模拟真实使用场景。import torch from torchvision import transforms, models from PIL import Image device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsNone) num_ftrs model.fc.in_features model.fc torch.nn.Linear(num_ftrs, 10) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.to(device) model.eval() class_names [bread, cheese, ice_cream, apple, banana, corn, carrot, fried_egg, rice, spaghetti] 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]) ]) img Image.open(test_food.jpg).convert(RGB) img_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(img_tensor) probabilities torch.softmax(outputs, dim1) top_prob, top_class torch.topk(probabilities, k3) for i in range(top_prob.size(1)): idx top_class[0][i].item() print(f{class_names[idx]}: {top_prob[0][i].item():.4f})这里有个小坑就是Image.open之后最好强制 convert(RGB)。有些食物图片是 RGBA 四通道的 PNG如果不转的话输入张量就是 4 个通道和模型期望的 3 个通道不匹配直接报错。模型 state_dict 里只保存了参数没有保存结构。加载模型时要先创建一个相同的模型结构再 load_state_dict所以我在推理脚本里重新定义了模型。还有一种更省事的方式是保存整个模型对象但不推荐因为模型结构在 PyTorch 升级后可能不兼容。6.2 混淆矩阵分析跑完测试集之后我输出了一张混淆矩阵图我就不贴了重点看几个数据点。意面带最容易和玉米混淆。细看原因是两者的颜色接近都是黄棕色系而且食物照片中都有长条形结构。模型没能学到“意面是弯曲细长的面食玉米是一粒一粒的”这种语义级区分。香蕉和胡萝卜也非常容易分不清。测试集里胡萝卜图片有不少是切成条状的和香蕉外形很接近。这提醒了我真实分类场景中“形状”和“颜色”并不总是可靠特征模型如果能捕捉到“香蕉有弧度”“胡萝卜根部有渐细的尖端”效果会更好。通过混淆矩阵我找到了 3 个样本量不足的难分对有针对性地多收集了这些类别的图片加入训练集。下次迭代的验证准确率从 85% 提升到了 88%。6.3 模型大小与推理性能训练完成之后我对模型的工程指标做了记录。指标数值模型参数量约 1120 万模型文件大小约 43MB单张图片 CPU 推理耗时约 45ms单张图片 GPU 推理耗时约 3ms显存占用约 1.2GB对于移动端或嵌入式部署1120 万参数量还是偏大。项目如果后续扩展我会考虑蒸馏成一个小型网络或者直接用 MobileNetV3。但在“验证思路、跑通流程”这个阶段ResNet18 的体量是完全合理的。7. 经验的总结做图像分类项目最容易忽略的几件事整个项目跑下来我自己有几点特别深的体会。数据整理阶段多花一小时训练阶段省一周。目录结构、类别命名、数据校验、样本分布检查这些不起眼的准备决定了后面所有环节是否顺畅。我一开始没有认真检查每个类别的图片数量训练到中途才发现有一个类别图片数量比其他类别少了一半导致模型严重偏向样本多的类别又要重新整理数据从头训练。不要迷信模型越复杂越好。这个 10 分类任务用 ResNet18 已经溢出真正决定模型上限的是数据质量、增强策略和调参细节。换更大的模型不会带来想象中的大提升。手写训练循环是值得的。虽然 PyTorch Lightning 之类的框架能省掉大量样板代码但对我来说手写循环带来的收益是理解层面的——理解了每个张量的变化、每行代码的作用之后后面用任何高阶框架都只是一个“选择”的问题而不是被框架牵着走。保存模型时多保存一份带优化器状态的完整检查点。我后来想接着上次的训练继续跑发现只保存了模型参数优化器的动量状态没保存结果后面几个 epoch 的收敛速度明显变慢相当于丢掉了一部分训练信息。正确做法是每一轮同时保存 optimizer.state_dict 和调度器状态torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_val_acc: best_val_acc, }, checkpoint.pth)恢复训练时就三行代码搞定checkpoint torch.load(checkpoint.pth) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict])最后再说一个小技巧是我后来才养成的习惯每训练一个模型就在最后带着 test_loader 跑一遍混淆矩阵同时把预测错误的样本图片存下来按类别分文件夹放好。这样做有两个直接好处一是能一目了然知道困难样本长什么样二是给后续的数据增强方向提供了明确依据。前几次做项目我就是只盯着准确率遇到瓶颈根本没有方向直到养成这个习惯之后模型的迭代才真正开始“有路线图”。整个食物分类项目从头到尾一共花了不到两周的下班时间却把 PyTorch 的完整工作流彻底打通了。如果你也是刚入门有一定 Python 基础想找一个有真实场景又不至于复杂的项目食物图片分类是个不错的选择。照着这个思路做一遍动手敲一遍代码调试一遍报错你会发现自己对深度学习模型训练的理解一下子变得具体了很多。
返回列表
PREV
查看更多资讯
NEXT
返回资讯列表