资讯动态

PyTorch图像分类实战:3小时跑通模型训练全流程

发布时间:2026/10/19 10:11:37 来源:尧图企业网站定制
先坦白一句标题里的“3小时”不是玄学是我把很多新手的真实进度压缩之后的时间盒。如果你也在PyTorch图像分类这条路上卡了好几天脚本抄来抄去却总在数据加载、维度对不上、显卡报错这些地方打转那这篇就是给你准备的。我会按一条我验证过很多次的路线带你从数据目录、数据集加载、模型搭建、训练循环走到评估优化和模型保存把PyTorch图像分类的全流程完整踩一遍并把新手最容易翻车的位置提前标出来。这条路线默认你已经能跑起Python和PyTorch不要求你懂太多理论但要求你愿意照着敲代码、愿意在报错里多看一眼堆栈。适合的人群是刚看完深度学习基础视频、想用一个完整项目把“数据加载—训练—评估”串起来的人也适合写过一些分类脚本但总觉得流程缺一块的人。3小时能不能跑完关键在于你愿不愿意先砍掉花哨的部分只留主干。1. 3小时不是吹的先把全流程拆成五个里程碑为什么很多新手花一周都跑不通一个图像分类项目我观察到的原因不是大家笨而是大家想一步到位一边学数据处理、一边抠网络结构、一边研究各种论文里的训练技巧最后被自己的代码量淹没了。所以这次我建议你换个思路——把整个项目拆成五个小里程碑每个里程碑只有一件事做完就进入下一步绝不在某一个环节里反复优化。1.1 新手的第一个瓶颈往往不在模型你能猜到一个分类项目里新手花时间最多的地方是哪里吗是训练脚本里的模型定义不是模型反而是最好抄的部分。真正耗时间的是三类问题数据加载不顺畅、显卡报错、以及训练跑起来之后看不懂结果。这三类问题恰好都出现在“模型之外”所以我的第一个建议是不要一上来就纠结要不要自己写一个DenseNet或者ConvNeXt先用最标准的ResNet骨架跑通闭环建立对整个流程的信心。另一个常见的错觉是“我必须把代码写得非常通用、非常优雅”。如果你面向的是学习或一个快速验证性的项目完全没必要。比如数据读取不一定要写一个继承自Dataset的子类torchvision提供的ImageFolder大多数场景够用训练函数不一定要支持断点续训、多卡并行、混合精度全开先跑出一个能稳住下降的Loss再谈花活。这叫最小可用闭环。1.2 我的时间盒规划80分钟给训练20分钟给排坑所谓时间盒就是把总时间切给几个固定阶段到点就进入下一步不恋战。下面是我用下来最顺的分配表时间区间里程碑完成标准0-30分钟数据准备与检查成功打印出训练集、验证集的大小和类别数能看到一个batch图像的shape30-60分钟模型加载与训练骨架训练脚本能跑起来哪怕Loss还很差至少不报错60-100分钟第一轮有意义的训练Loss明显下降验证集准确率超过随机猜测100-140分钟观察与针对性优化学率率、数据增强做一轮调整曲线更平稳140-180分钟评估、可视化与保存模型输出混淆矩阵和分类报告保存最优模型注意每个阶段的“完成标准”都是可验证的不是感觉上的“差不多了”。我见过太多人卡在数据准备阶段反复清洗数据、反复看盘数据分布60分钟过去了还在打转。清洗数据很重要但不是今天这个3小时场景的核心目标跑通闭环才是。1.3 三个前提条件缺一个都别谈3小时第一你最好有块能用的GPU。哪怕是很老的显卡只要能跑CUDA速度也比CPU快一个量级。没有GPU也不是不能学只是3小时会变成6小时而且你会被“等待”消磨掉耐心。第二数据已经在磁盘上并且已经按类别分好文件夹不要在3小时里同时完成数据爬取和清洗。第三使用迁移学习也就是直接加载别人在大规模数据集上预训练好的权重你只替换最后分类层去适配你自己的类别数。三条缺一风险都会大幅上升。我把迁移学习放到第三个前提是因为它真的能救新手。从零训练一个十层以上的网络在没有充分的调参经验时大概率遇到收敛慢、过拟合、Loss乱跳等问题。而预训练模型已经把边缘、纹理、形状这些低级特征学好了你的任务只是让它在你的小数据集上做“最后一公里”的适配。2. 数据准备这半小时目录、Transform和DataLoader的接线数据准备听起来无聊却是我见过翻车率最高的环节。很多人把目录结构搞乱、归一化参数抄错、训练集和验证集用同一套数据增强最后训练曲线看着正常验证准确率却上不去。所以这里我建议多花5分钟把地基打稳。2.1 ImageFolder的标准目录结构别自己发明torchvision的datasets.ImageFolder对目录结构有明确要求根目录下每个子文件夹代表一个类别子文件夹里的图片都属于这个类别。一个典型结构是data/ ├── train/ │ ├── 类别A/ │ │ ├── img1.jpg │ │ └── img2.jpg │ └── 类别B/ │ └── img3.jpg └── val/ ├── 类别A/ └── 类别B/实践里最容易犯的错是类别文件夹名字带中文或空格或者类别A和类别B的文件混放在同一个文件夹再靠一个csv文件去标注。ImageFolder不支持这种玩法所以要么你花时间重排目录要么就得自己写Dataset。对新手来说我强烈建议多花几分钟重排目录因为ImageFolder能免费帮你完成三件事自动扫描类别、自动给类别编号、自动建立索引。这三件事自己写很容易出bug。2.2 训练和验证一定要用两套TransformTransform就是你对图像做的预处理核心包括缩放尺寸、裁剪、翻转、转张量、归一化。我直接给你一套可以放心抄的配置前提是你后面用的预训练模型是ImageNet系列的。训练集用随机尺寸裁剪和随机水平翻转验证集用等比缩放加中心裁剪。为什么要分开因为训练集的随机增强是为了让模型看到更多样子的同一种对象减少过拟合验证集要稳定地反映模型真实水平任何随机操作都只会带来波动。from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), 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]) ]) train_dataset datasets.ImageFolder(data/train, transformtrain_transform) val_dataset datasets.ImageFolder(data/val, transformval_transform) print(len(train_dataset), len(val_dataset)) print(train_dataset.classes)这里面有两个细节值得解释。第一个是Resize(256)再CenterCrop(224)的组合把短边统一到256再取中心224x224的区域这样既保证了输入尺寸一致又不会像直接Resize(224)那样把长宽比压坏。第二个是Normalize里的mean和std这是ImageNet数据集的统计值迁移学习中大家默认都用这一组。如果你不是用ImageNet预训练权重这个值需要重新统计但新手阶段直接沿用这一组没有大问题。另一个高频坑是有人把train_transform整套照搬到val_transform上。结果验证时图像也被随机裁剪和翻转导致每次评估的结果差别很大无法判断模型到底有没有进步。2.3 DataLoader里关系到“能否3小时跑完”的线程参数DataLoader是负责批次加载的组件它有几个参数会直接决定你GPU是不是在空转。第一个是num_workers表示用几个子进程做图像解码和预处理。取0表示在主进程里做GPU经常要等待CPU训练速度慢一大截取4或者8通常能明显缓解。第二个是pin_memory设为True会把数据放进锁页内存GPU取数据更快前提是操作系统和内存足够。第三个是shuffle训练集必须True验证集必须False。验证集打乱没有任何意义还会让你画混淆矩阵时对不上顺序。train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue)如果你在跑实验时发现GPU利用率一直很低先别急着查代码逻辑大概率是num_workers太小或者每个batch里有图片解密耗时太长。3小时的项目里这一项是“时间成本放大器”值得你早点重视起来。2.4 动手检查数据形状和类别数量写DataLoader之后我建议你跑一段最小的检查代码确认张量形状完全符合预期。这件事只需要15秒却能省掉后面模型里令人抓狂的维度报错。import torch images, labels next(iter(train_loader)) print(images.shape) # 期望 [batch_size, 3, 224, 224] print(labels.shape) # 期望 [batch_size] print(torch.max(labels), torch.min(labels)) # 类别编号从0开始如果这里输出的形状不是[64, 3, 224, 224]那就别急着碰模型先回去查Resize和Crop的尺寸。另外确认一下labels是连续整数编号后面用CrossEntropyLoss时这是硬性要求。3. 模型环节为什么我劝你别自己写网络直接上迁移学习到了模型这步我给新手的建议和很多人想的不一样。大家总觉得“图像分类的灵魂是设计一个好的网络结构”但那是做研究和发论文的角度。对于快速出成果、验证流程最重要的不是结构创新而是“用一个已经被验证过足够强大的网络快速适配你的数据”。3.1 两行代码加载ResNet预训练模型PyTorch的torchvision.models里集成了很多图像分类模型。以ResNet18为例加载预训练权重的现代写法是这样import torchvision.models as models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT)注意很多旧教程里写的pretrainedTrue在新版本PyTorch中会弹出弃用警告我建议直接使用weights参数。ResNet18虽然只有约1100万参数但作为迁移学习的主干网络已经很够用尤其适合数据量不大的项目。如果你的显卡比较好想追求更高准确率可以把resnet18换成resnet50代码只改一个名字。加载之后要替换最后的全连接层因为原来这个模型是在ImageNet的1000个类别上训练的你的数据集类别数量通常不是1000。每个模型的最后一层名字不同ResNet系列叫fcEfficientNet系列叫classifier我用代码先确认再替换in_features model.fc.in_features num_classes len(train_dataset.classes) model.fc torch.nn.Linear(in_features, num_classes) print(model.fc)这里新手容易犯的错是不知道in_features怎么来。其实通过model.fc.in_features就能自动拿到上一层的输出维度不需要手动算。替换之后分类层的输出数量就等于你的类别数。3.2 冻结前层还是微调全模型取决于你的数据量加载预训练模型之后有两种常见策略。第一种是全量微调也就是所有层都继续更新梯度。第二种是冻结住特征提取层只训练自己添加的分类层。到底选哪种我的判断逻辑很简单如果每个类只有几百张图数据量不大那分类层的随机初始化和预训练特征的分布之间会有不小的差距此时只有分类层能训练很考验学习率和迭代次数。因此我通常建议先用“冻结除最后两三个层之外的所有层”的折中方案既保留预训练特征又让高层的语义特征有机会适配你的数据。实际代码里冻结层的方式是设requires_grad。比如我想冻结所有层只解开fc层for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True被冻结的层在优化器传参时就不要放进去了否则它不更新但你还给它计算梯度白费时间。更干净的写法是把需要更新的参数单独筛出来optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3 )很多人在冻结层之后忘了修改优化器参数的来源导致模型一直在“假装微调”Loss也在降但其实是靠分类层在硬扛整体收敛速度很慢。3.3 优化器、损失函数和调度器的初步搭配图像分类最常用的损失函数是CrossEntropyLoss这个没有太多需要犹豫的。优化器我建议新手直接上AdamW它对学习率的敏感程度比SGD低更容易稳定跑起来。如果你想让上限更高后面的优化阶段可以切到带Momentum的SGD配合余弦退火调度器。但第一轮训练别用SGD别问为什么等你有心情去调试Initial learning rate和momentum的组合时再说。学习率调度器我习惯用ReduceLROnPlateau它会在某个指标比如验证Loss停滞时自动把学习率乘一个系数。这样即便初始学习率不太合适后面也有自救的机会import torch criterion torch.nn.CrossEntropyLoss() optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3 ) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.1, patience2 )这段代码的意图很直接训练几步之后如果验证Loss不再下降学习率自动缩小到原来的十分之一给模型一次更精细的调整机会。我遇到过不少新手从头到尾用一个固定学习率最后模型卡在一个局部最优附近怎么也动不了就是这个原因。3.4 显存不够时先查这几处而不是换显卡3小时的项目里最怕遇到CUDA out of memory。新手第一反应可能是“我显卡太差了”但实际上绝大多数OOM不是硬件不够而是有几个小地方浪费显存。第一输入图像尺寸是否设成了384甚至512对新手任务来说224完全够用。第二batch_size是不是设得太大64跑不动就降到32、16准确率不会因为你batch_size大而突变。第三有没有用AMP半精度训练下面一节我会给模板这一步通常能省近一半显存。如果这些都调完了还是OOM再检查是不是你在验证时也保存了整批梯度。验证阶段用torch.no_grad()包住就能避免梯度占显存别小看这个操作很多人忘记。4. 训练循环与监控从能跑出数字到看得懂曲线有了数据和模型下一步就是训练循环。这一部分我不推荐你直接复制那些上百行的训练框架而是先写一个最朴素的版本理解每个关键动作的意义然后再逐步加功能。4.1 一个最小可运行的训练循环逐行解释哪里会踩坑训练一个epoch的骨架长这样def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_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() total_loss loss.item() * images.size(0) _, preds outputs.max(dim1) correct preds.eq(labels).sum().item() total labels.size(0) avg_loss total_loss / total acc correct / total return avg_loss, acc这段代码里有几个坑几乎是新手必踩。第一个是忘了model.train()或model.eval()。model.train()会启用Dropout和BatchNorm的统计更新而验证时必须切到model.eval()否则同一个batch的数据进来、同样的权重输出结果却会因为Dropout随机性而变动导致验证指标不准。第二个是optimizer.zero_grad()必须放在前向之前。放在backward之后、下一个batch前也不是完全不行但更容易漏漏一次梯度就会累加非常难排查不如规规矩矩放在loss计算之前。第三个是outputs.max(dim1)取预测类别时dim1表示按类别维度取最大值索引写成dim0就会得到完全不靠谱的结果。验证函数反过来写没有梯度计算没有optimizer只为了评估torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() total_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) total_loss loss.item() * images.size(0) _, preds outputs.max(dim1) correct preds.eq(labels).sum().item() total labels.size(0) return total_loss / total, correct / total注意这里的torch.no_grad()装饰器它能让验证过程完全不保存计算图省显存、提速还避免验证步骤把梯度混进模型。4.2 AMP半精度训练不动模型结构直接省一半显存如果你的显卡支持AMP是性价比最高的一项改动。核心思路是前向推理用半精度浮点数计算反向传播前用梯度缩放器避免精度下溢。PyTorch的写法很固定scaler torch.cuda.amp.GradScaler() for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()对比普通训练改动只有三行创建一个GradScaler用autocast包住前向和loss计算把loss.scale之后再backward。我见过有人把scaler.step(optimizer)和scaler.update()写反结果训练Loss一直不降。顺序是固定的先scaler.scale(loss).backward()再scaler.step(optimizer)最后scaler.update()。三行一起出现才是完整循环。AMP不是只在高端卡上有用很多消费级显卡也能跑。建议你第一轮就用上这样即使batch_size设得稍大也不用担心OOM打断节奏。4.3 怎么监控训练曲线tqdm和TensorBoard3小时的项目每个epoch都要能看到训练进度否则你根本不知道是自己代码卡住了还是GPU在思考。最省事的是tqdmfrom tqdm import tqdm for epoch in range(epochs): model.train() pbar tqdm(train_loader, descfEpoch {epoch1}/{epochs}) for images, labels in pbar: # 训练代码... pbar.set_postfix(lossloss.item())如果你希望训练完之后能翻看曲线用TensorBoard是最合适的。在训练循环里写两行from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/exp1) # 在每个epoch结束后 writer.add_scalar(train/loss, train_loss, epoch) writer.add_scalar(val/loss, val_loss, epoch) writer.add_scalar(val/acc, val_acc, epoch)然后终端里执行tensorboard --logdir runs浏览器打开本地端口就能看到Loss和准确率曲线。注意几点不要每个step都写写epoch层面的数据就够了目录名不要用默认的runs/否则多个实验混在一起很难看TensorBoard里的曲线比你打印在终端里的数字更能暴露问题尤其是震荡和过拟合。4.4 训练正常但验证不涨这是最常见的调试场景当你看到训练Loss在下降、训练准确率接近100%但验证准确率卡在低点或者大幅震荡时不要慌。按这个顺序排查首先确认验证集没有做和训练集一样的增强随机裁剪会直接导致验证结果不稳定其次怀疑过拟合数据量小时模型会死记训练集解决方法是加数据增强、减小模型容量、或者把早停策略加进来最后检查你的数据是否真的按类别均分成了两份切片时如果用了shuffleFalse再按固定比例切很容易造成某个类别在验证集里只有十几张指标波动自然大。另外验证集准确率比训练集还要高也是一种异常信号。这种情况通常不是模型厉害而是数据划分泄漏比如验证集里混进了和训练集完全相同的图片或者两个集合来源不同导致验证集更简单。处理的办法是重新划分数据并保证两个集合的分布来源尽可能一致。5. 模型评估与错误样本分析别只盯着accuracy跑完训练很多人拿到一个准确率就认为项目结束了。这个习惯对新手尤其危险因为最终准确率只能告诉你“整体好不好”不能告诉你“哪里不好”。我强烈建议你多花10分钟做三个动作看分类报告、画混淆矩阵、把预测错的样本可视化。5.1 用sklearn输出classification_report和混淆矩阵分类报告能按类别给出精确率、召回率、F1值。这在类别不平衡时特别重要——总准确率90%听起来很好但某个类别的召回率可能只有0.3相当于这个类几乎总被认错。代码很简单from sklearn.metrics import classification_report, confusion_matrix import numpy as np all_preds, all_labels [], [] model.eval() with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) preds outputs.argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_namesval_dataset.classes)) print(confusion_matrix(all_labels, all_preds))注意这里又出现了一次model.eval()因为评估阶段绝不能启动Dropout和BatchNorm那些随机机制。另外我特意把预测结果先移到CPU再转numpy避免GPU和CPU的数据类型问题。怎么看混淆矩阵呢对角线上的数字越大越好矩阵里最大的非对角线数字就是最容易混淆的类别对。比如A类有200张图其中70张被预测成了B类那你就知道这两个类别在视觉上可能很相似或者你的模型对这两个类别学到的特征区分度不足。5.2 把模型“想错了”的样本画出来数字不够直观直接把错误样本画出来才是提升游戏理解的最好方式。思路是在验证集里收集所有预测错误的样本记录图像、真实标签、预测标签然后画成网格。import matplotlib.pyplot as plt import torchvision.transforms.functional as F misclassified [] for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) preds outputs.argmax(dim1) for i in range(images.size(0)): if preds[i] ! labels[i]: misclassified.append((images[i].cpu(), labels[i].item(), preds[i].item())) if len(misclassified) 16: break if len(misclassified) 16: break fig, axes plt.subplots(4, 4, figsize(12, 12)) for idx, (img, label, pred) in enumerate(misclassified): row, col divmod(idx, 4) # 反归一化img是经过Normalize的要先恢复成可视化可用的范围 img img * torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) img img torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) img img.clamp(0, 1) axes[row][col].imshow(F.to_pil_image(img)) axes[row][col].set_title(f真:{val_dataset.classes[label]} 预:{val_dataset.classes[pred]}) axes[row][col].axis(off) plt.tight_layout() plt.show()这里有一个反归一化的操作我特别标出来了。训练数据经过Normalize之后像素值不再是0到1的正常范围直接imshow会看到发灰发黄的图片要么看不懂要么误以为数据坏了。恢复方法就是乘标准差再加均值然后clamp到0和1之间。5.3 从错误样本里反向发现数据问题错误样本可视化不只是为了截图发朋友圈它的核心价值是帮你反向检查数据质量。我做过一个宠物分类项目模型把所有黑猫都预测成了狗。查看错误样本后发现那一批黑猫照片里有的背景占了很大面积模型明显在学背景而不是学猫本身。这就说明数据本身有很大噪声单靠调模型参数是补不回来的要去做数据清洗或者换更干净的子集。另一个常见现象是两个类别之间本身就存在重叠比如“林间土路”和“公路”很多图的标签其实是一个人的主观判断。这种问题模型很难解决你需要回到标注层面统一标准。简单说错误样本分析能把“模型的锅”和“数据的锅”分开这是新手最缺少的判断力。5.4 保存checkpoint的正确姿势别只存state_dict训练结束时你可能想保存模型权重很多人一上来就是torch.save(model.state_dict(), model.pth)。这样不是不能用但如果你后面想继续训练、调学习率或者对比实验却发现优化器状态和epoch都丢了体验会很痛苦。我建议从第一天就养成保存完整checkpoint的习惯best_acc 0.0 for epoch in range(epochs): # 训练和验证... if val_acc best_acc: best_acc val_acc torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_acc: best_acc, }, best_checkpoint.pth) print(f已保存最佳模型验证准确率: {best_acc:.4f})加载的时候反过来就行checkpoint torch.load(best_checkpoint.pth, map_locationdevice) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) epoch checkpoint[epoch]用map_location指定设备能避免你在GPU机器上训练、后来拷到没GPU的机器上加载时报错。还有一个小经验如果一个训练过程需要好几天我会同时保存“最近一个epoch”和“验证集最佳”两份checkpoint后者用于最终评估前者用于一旦训练中断能继续跑。6. 新手的坑清单从报错到玄学问题的排查链路写到这你已经能跑通一个完整的图像分类项目了。但这毕竟是一条由很多报错串起来的路我把新手最常见的几类坑按排查顺序整理成清单遇到问题先对号入座。6.1 环境类的坑OOM、worker崩溃、数据加载卡死CUDA out of memory的处理思路前面已经说过降低batch_size、确认224尺寸、开AMP、验证阶段加no_grad。如果这几个都调完还OOM那就需要看看你的代码是不是在同一个GPU上同时跑着别的程序。DataLoader的worker崩溃一般表现为训练刚开始就报错而且报错信息里带fault或者segmentation。大多数情况是num_workers设得太大或者数据文件里有损坏图片。解决办法先把num_workers调成0跑一遍确认是不是图片损坏如果0能跑再逐步调高num_workers。数据加载卡死常见于Windows环境把pin_memory和num_workers这两个参数组合改成num_workers0、pin_memoryFalse就能绕过去虽然慢一点但至少稳定。6.2 数据类的坑类别不平衡、增强过度、归一化写错类别不平衡时模型会倾向把什么都预测成样本数量最多的类别最终准确率可能挺高但少数类几乎全错。看classification_report能立刻发现这个问题。最省事的处理方式是加权损失给CrossEntropyLoss传一个weight让样本少的类别受到更大惩罚。class_counts torch.bincount(torch.tensor(all_labels)) class_weights 1.0 / class_counts.float() class_weights class_weights / class_weights.sum() criterion torch.nn.CrossEntropyLoss(weightclass_weights.to(device))数据增强做得太狠也会出问题。记住一个原则增强是模拟真实世界里对象的各种形态不是把图片弄到人类都认不出来。如果你在训练集上把旋转角度拉到45度以上、亮度抖动幅度拉到0.5模型的学习负担会暴增Loss反而下不去。归一化写错的情况是很多人把mean和std的顺序搞反了或者用成了[0.5, 0.5, 0.5]。对于ImageNet预训练模型请统一使用mean[0.485, 0.456, 0.406]、std[0.229, 0.224, 0.225]并且在训练和验证两套transform里保持一致。不一致会导致预测结果莫名其妙地差几个百分点。6.3 训练类的坑Loss为NaN、模型不收敛、结果不可复现Loss变成NaN优先查三件事输入数据是否包含NaN像素、学习率是不是过大、标签是否是连续的0到类别数减一。三者里最隐蔽的是最后一个——如果标签里有-1或者100这种越界值CrossEntropyLoss计算时会直接NaN。排查办法是在数据检查阶段加一句断言确保labels.max()小于num_classes。结果不可复现让很多新手抓狂。昨天跑出85%准确率今天一模一样的脚本变成83%不是玄学是随机性没有被控制住。在训练脚本开头固定几个种子import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)需要说明的是即使固定了这些种子DataLoader多进程处理时的随机性在某些平台下依然不能完全消除。如果遇到这种情况把num_workers设为0再复现一次通常就能对上。6.4 一套推荐的排错顺序表真到报错时按顺序排查比胡乱搜效率高很多。我给自己定了一套顺序分享给你排查顺序检查对象常用手段1数据和标签链路打印一个batch的shape、标签范围、类别数2模型输入输出确认in_features、out_features、device一致3训练循环结构确认train/eval、zero_grad、no_grad位置4Loss表现观察第一个epoch的loss是否在一个合理初始范围5硬件资源GPU利用率、显存占用、数据加载瓶颈6随机性固定种子排除不可复现干扰这套顺序最大的价值是逼你先确认最简单的环节再碰玄学问题。很多时候报错信息指向模型但真正的原因是数据和标签不匹配。最后说几句实在的照这条路线走完你手里的东西已经不是一个练习脚本而是一套完整的图像分类工作流数据组织、迁移学习、训练监控、评估分析、模型保存每一环都有可复现的代码和判断标准。我自己的体会是3小时真正要练的不是“把所有代码背下来”而是学会在每一个里程碑里快速判断“下一步该做什么”。遇到报错先看堆、再看数据、再看模型结构而不是盲目改参数。如果第一遍跑不完也不要硬扛把某个里程碑切出来重跑一次往往第二次就顺畅了。最后分享一个小技巧项目做到后半段我会把epoch数先设成3用最小规模的训练去验证整条链路完全通畅确认无误后再改成正式epoch数跑长时间训练。这个“小跑热身”的做法帮我省掉的等待时间远比想象中多。希望你也能用这条路线在3小时内完成第一个真正意义上的PyTorch图像分类项目。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑