资讯详情

资讯详情

PyTorch训练脚本全流程详解:从数据加载到模型保存的标准套路

1. 训练脚本里到底在做什么先梳理一遍完整流程写过几行模型代码的初学者往往卡在一个很尴尬的阶段看文档能知道Conv2d是卷积层知道ReLU是激活函数模型定义也能照着抄但一个完整、能跑、还能装上验证逻辑的训练脚本就算拼出来也总觉得哪里别扭。有的代码把训练和验证混在一起有的不知道在哪个位置换model.eval()有的干脆把loss.item()和loss的区别丢在一边直接print(loss)精度曲线一塌糊涂。说到底大家缺的不是某一个API的用法而是一条贯穿始终的“训练套路”。1.1 一次训练过程的全景图从数据到参数更新任何监督学习任务训练脚本本质上都在重复同一件事把一批数据送进模型让模型给出预测计算预测和真实标签之间的差距然后根据差距反向传播更新参数。这个过程分成四个阶段数据准备。把原始数据变成张量按固定大小切块分批喂给模型。前向计算。数据流过模型的每一层最后输出一个预测值。反向传播。自动求导机制算出每个参数对损失的梯度也就是每个参数应该往哪个方向调、调多少。参数更新。优化器拿着梯度去修改模型参数让损失变小一点。这四个动作看起来简单但真正写成脚本时很多人会栽在第1步和第3步上。数据张量到底是NCHW还是NHWC梯度要不要先清零loss能不能直接累加这些细节如果不理解代码连报错都看不懂。1.2 三个关键角色数据集、模型、优化器以及它们怎么配合训练套路里有三个核心对象它们各司其职Dataset负责根据索引取出一条样本DataLoader负责把这些样本按批次打包并在训练时打乱顺序。Model封装了网络的层结构和前向计算方式同时保存了所有可学习参数。Optimizer负责拿着梯度去更新模型参数。它不参与前向计算只关心“模型有哪些参数”和“当前梯度是多少”。这三者在代码里并不是偶发的协作关系。数据集和模型之间靠张量形状对接模型和优化器之间靠model.parameters()对接。只要这两个接口形状对上了训练脚本的其余部分基本就是一成不变的循环结构。这也正是“套路”这个词的意义你不需要每次重新设计训练流程只需要按这条链路套用即可。2. 训练循环的骨架代码每个环节为什么这么写明确训练的整体流程后下一步是把流程固化成代码。我给出一套最通用的骨架然后再逐行解释每个部分背后的设计逻辑。2.1 环境准备与设备判断别一开始就把device写死很多人习惯直接写device torch.device(cuda)这在有显卡的机器上没问题可一旦换到没显卡的环境下一秒就是报错。更稳妥的做法是先做设备检测import torch device torch.device(cuda if torch.cuda.is_available() else cpu) print(Using device:, device)这样做的好处很明显同一份脚本既能在本机的CPU上调试也能在带卡的服务器上跑出更高效率。调试阶段我建议直接用CPU跑通小规模数据确认逻辑没问题再切GPU不然一上来就带着显存报错排查问题成本会高很多。2.2 DataLoader的迭代逻辑为什么训练要按batch来训练不是把整个数据集一次性塞进模型。一方面数据量太大显卡内存放不下另一方面分批训练时的梯度更新带有一定的随机性这种随机性反而能帮助模型跳出局部最优点。PyTorch的DataLoader就是干这件事的from torch.utils.data import DataLoader from torchvision import datasets, transforms transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) train_dataset datasets.CIFAR10( root./data, trainTrue, downloadTrue, transformtransform ) train_loader DataLoader( train_dataset, batch_size64, shuffleTrue, num_workers2 )shuffleTrue是训练集的关键配置。如果不打乱顺序模型会按固定的样本顺序反复学习容易记住序列本身而非数据特征导致训练早期收敛变慢甚至驯出一个在顺序上“作弊”的模型。num_workers是加载数据的子进程数量在Windows上有时会引发出入队错误如果遇到这类问题可以临时改为0先跑通再说。2.3 前向、反向、更新的标准三连zero_grad到底在清什么每个批次内的核心代码就是下面这几行几乎所有PyTorch训练脚本里都能看到它的变体optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step()新手最容易忽略的是optimizer.zero_grad()。PyTorch的梯度是累积的它会不断往.grad属性里累加新值而不是在下一次backward前自动清空。如果不手动清零每个batch的梯度就会叠加上前一个batch的梯度参数更新方向越来越歪损失曲线通常表现为震荡幅度特别大且不收敛。loss.backward()做的是反向传播它真正干的活是从loss出发沿着计算图往回走给每个参与运算的张量计算梯度并把结果写进对应张量的.grad属性里。注意这个步骤之后参数并没有被修改真正的修改发生在optimizer.step()。一前一后职责非常清晰。2.4 一个可直接运行的完整示例CIFAR-10上的简单分类器把上面的片段拼成一个完整可执行的脚本这是训练套路最标准的模板import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) train_dataset datasets.CIFAR10( root./data, trainTrue, downloadTrue, transformtransform ) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 32, kernel_size3, padding1) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.fc1 nn.Linear(64 * 8 * 8, 256) self.fc2 nn.Linear(256, 10) def forward(self, x): x torch.relu(self.conv1(x)) x torch.max_pool2d(x, 2) x torch.relu(self.conv2(x)) x torch.max_pool2d(x, 2) x x.view(x.size(0), -1) x torch.relu(self.fc1(x)) x self.fc2(x) return x model SimpleCNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) num_epochs 10 for epoch in range(num_epochs): model.train() running_loss 0.0 for inputs, labels in tqdm(train_loader): inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) epoch_loss running_loss / len(train_dataset) print(fEpoch {epoch1}/{num_epochs}, Loss: {epoch_loss:.4f})这段代码里的两处设计值得一提。第一running_loss累加的是loss.item()而不是loss。loss本身是一个带计算图的张量直接累加会把整个反向图保留下来内存占用越叠越高几十个epoch后很容易把内存撑爆。loss.item()取出的是标量准确说是一个Python浮点数累加的就是纯数值。第二每个epoch结束时除以len(train_dataset)而不是len(train_loader)得到的是整个数据集的平均损失不是某个批次的损失。这样打出来的曲线明显平滑得多。3. 用套路承接变化验证、保存与断点续训训练循环的骨架足够稳定但真正用于实际项目时只跑训练循环是不够的。你需要知道模型在没见过的数据上表现怎么样需要随时把当前最好的权重留一份下来甚至中断后还能接着跑不用全部从头再来。3.1 model.train() 与 model.eval()两行代码背后的行为差异我在很多新手代码里看到过两种极端情况要么从来不写model.train()和model.eval()要么每次训练和验证都只写model.eval()忘了切换回train()。这两行代码在模型结构里不是摆设它们会直接影响部分层的行为。model.train()模式下Dropout是开启的BatchNorm在逐batch计算当前批次的均值和方差。model.eval()模式下Dropout直接不工作BatchNorm改为使用训练期间累积的全局统计量。如果验证时代码还停留在train()状态Dropout会让每次验证的输出带有随机性同一份数据跑两次结果不一样评估出的精度指标根本不稳定。反之如果训练时忘了从eval()切回train()Dropout全程关闭模型相当于一个精简版本表达能力不足精度的天花板就打不上去。所以请务必记住这个铁律训练循环前写model.train()验证循环前写model.eval()。3.2 验证集评估的正确打开方式no_grad与对整个epoch做平均验证阶段和训练阶段最大的区别是验证过程不需要更新参数因此也完全不需要为反向传播保留梯度。torch.no_grad()就是用来做这件事的上下文管理器它让PyTorch在前向计算过程中不构建计算图内存占用更低、速度也更快。model.eval() val_loss 0.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) val_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_loss / len(val_dataset) val_acc correct / total print(fVal Loss: {val_loss:.4f}, Acc: {val_acc:.4f})这里有一个很多人踩过的坑验证时计算准确率用的是torch.max(outputs, 1)它返回两个东西——最大值和最大值所在索引。如果直接写torch.max(outputs, 1)然后predicted result拿到的其实是最大值本身而不是索引算出来的准确率自然是错的。正确写法是像上面这样用一个下划线或一个临时变量把最大值丢掉只保留索引。3.3 checkpoint不是简单调save权重、优化器状态、epoch都要存保存模型最简单的操作是torch.save(model.state_dict(), model.pth)。对训练套路而言这份权重文件能完成最基本的推理任务但如果你训练到第30个epoch突然被中断恢复时却发现只有权重没有优化器状态Adam里记录的历史动量全部丢失后续训练效果会打折扣还得重新调整学习率浪费不少时间。因此完整的checkpoint应该至少包含三样东西保存内容对应代码恢复用途模型权重model.state_dict()恢复网络结构内的可学习参数优化器状态optimizer.state_dict()恢复动量、学习率等内部状态训练进度epoch、best_acc得知从哪个epoch继续、当前最优成绩推荐的组织方式是把它们打包成一个字典集中存储checkpoint { epoch: epoch 1, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_acc: best_acc, loss: epoch_loss } torch.save(checkpoint, fcheckpoint_epoch{epoch1}.pth)3.4 恢复训练的常用代码恢复训练时大部分人会卡在“模型重载后找不到训练进度”这个环节。其实流程并不复杂checkpoint_dir ./checkpoints start_epoch 1 best_acc 0.0 import os if os.path.exists(checkpoint_dir): ckpt_files sorted(os.listdir(checkpoint_dir)) if ckpt_files: latest_ckpt os.path.join(checkpoint_dir, ckpt_files[-1]) ckpt torch.load(latest_ckpt, map_locationdevice) model.load_state_dict(ckpt[model_state_dict]) optimizer.load_state_dict(ckpt[optimizer_state_dict]) start_epoch ckpt[epoch] 1 best_acc ckpt[best_acc] print(fResumed from epoch {ckpt[epoch]}, best_acc {best_acc:.4f})恢复后进入训练主循环只需要把for epoch in range(num_epochs)改成for epoch in range(start_epoch, num_epochs 1)即可。另外需要注意的是如果你换了一张卡或换了机器torch.load时要指定map_location否则模型权重本来放在GPU上恢复时环境里没有对应设备就会报错。4. 训练过程中的异常信号与定位手段训练套路写好后真正折磨人的环节往往不是写代码而是面对各种奇怪的训练现象。很多新手看到一个Loss: nan就直接懵了四处搜代码改了一通参数结果完全没针对根因。这一节我总结几种最常见的异常信号以及排查方式。4.1 loss不降、梯度为0、loss为NaN三类高频问题先看表格对照再逐一展开现象可能原因快速排查方向loss不降学习率过大损失震荡不收敛打印梯度范数降低学习率loss一直不变模型输出恒定梯度消失检查是否有detach或no_grad误用loss变成NaN数据中存在NaN或梯度爆炸检查输入数据、梯度范数是否异常loss不降。训练前几个epoch损失纹丝不动先别怀疑模型结构十有八九是学习率出了问题。用一个从大到小的学习率试几组比如1e-2、1e-3、1e-4看哪一组能让loss先降起来。一般Adam优化器对1e-3这个默认值容忍度较高但如果你换了特殊损失或特殊网络这个值可能偏大。loss不变更像是恒等式。如果模型输出始终是同一个值损失自然恒定。出现这种情况常见原因是冻结了某些层requires_gradFalse或误把整个前向放在torch.no_grad()里。还有一种隐蔽场景是model.eval()之后忘了切回train()Dropout全程关闭模型能力被阉割损失就会卡在一个平台。loss为NaN。训练到一半突然变成NaN最直接的原因是梯度爆炸。你可以临时在loss.backward()后面加一句torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)把梯度限制在合理范围内。如果加上梯度裁剪后问题不再出现基本可以确定是梯度爆炸。另一种根源在数据侧如果你的标签里有NaN值或者输入图片有损坏导致的异常像素loss也会受影响这种情况需要在数据加载环节做过滤。4.2 过拟合与欠拟合的常见特征与对应策略训练损失和验证损失之间的差值是判断模型状态最直观的指标训练损失一直很高验证损失也很高欠拟合。说明模型容量不够、训练不充分或学习率过小。先加层数、加宽度或换更强的优化策略。训练损失持续下降验证损失不降反升过拟合。此时模型在背书而不是在学规律。应对策略是再加数据增强、加Dropout、减轻模型复杂度或提前停止。过拟合出现的时间点其实很有规律。大概在训练第10到第15个epoch之间验证损失会开始抬头而训练损失还在往下走。这时继续训练没有意义因为模型在验证集上的泛化能力已经到顶了。保存模型的逻辑也应该是“每轮验证时若精度高于历史最佳才保存”而不是“最后一个epoch结束才保存”。这样无论训练过程如何波动你手上始终有验证集表现最好的那一份权重。4.3 用print、tensorboard、简单日志定位训练问题训练套路本身是固定的但问题定位的手段可以很灵活。最简单的做法是在关键节点打点比如每个epoch打印完整体的loss和acc后再抽样打印两三个batch的梯度范数total_grad 0.0 for p in model.parameters(): if p.grad is not None: total_grad p.grad.norm().item() print(fEpoch {epoch1}, grad_norm: {total_grad:.4f})梯度范数如果小到1e-8量级说明梯度消失如果大到1e2以上说明梯度可能要爆炸。有了这个数值面对很多“loss不动”的诡异情况时就能快速判断出问题不是出在数据或模型结构而是出在反向传播这一层。再讲究一点可以把每个epoch的训练损失、验证损失、验证准确率、梯度范数按行追加到CSV文件里之后画图分析就有现成数据import csv with open(training_log.csv, a, newline) as f: writer csv.writer(f) writer.writerow([epoch 1, train_loss, val_loss, val_acc, grad_norm])这套东西组织好后后面每次训练新模型你都能快速回忆起“上一次在第几个epoch开始过拟合”“默认学习率会不会掉到NaN”而不需要重新翻代码猜参数。这也是“成套训练模板”带来的长期收益。5. 从一开始就值得养成的训练习惯训练套路跑通之后真正拉开差距的是细节习惯。同样的模型和数据集人家能每次复现同一份结果你换台机器结果就飘忽不定定位半天也找不到原因。习惯差异就在这里体现。5.1 随机种子怎么设深度学习里的随机来源很多模型参数的初始化顺序、数据打乱的排列方式、Dropout的随机选择。在相同代码下这些随机性会导致完全不同的训练结果这不是bug而是特性。如果你想让结果可复现最直接的办法是在脚本最前面统一设置随机种子import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42)这行代码的意义在于当你在自己的模型和别人的模型之间做对比时可以排除随机因素只观察结构差异带来的变化。很多开源项目的训练脚本开头都会有一段类似的代码不是玄学是严谨性的体现。5.2 学习率策略为什么也要当成套路的一部分lr1e-3从头跑到尾很多简单任务也能跑出不错的结果但对稍微复杂一点的任务这种僵化的策略并不划算。训练前期需要大学习率快速下降训练后期则需要小学习率谨慎收敛这就是学习率调度的目的。PyTorch里最简单的用法是每若干epoch将学习率乘一个衰减因子scheduler optim.lr_scheduler.StepLR(optimizer, step_size30, gamma0.1)记得每个epoch结束后调用scheduler.step()。如果你不理解它的作用也可以先照着这个句式用起来跑几次任务后你会自然体会到到了训练后半段固定学习率会在地形里来回“荡”而按计划衰减的学习率能在最后几个epoch把参数稳稳推进到更低的位置。5.3 记录训练信息模型名、参数、loss曲线的组织方式训练模型最好别只存一个裸权重文件。项目跑久了你会发现目录里堆了一堆model.pth、final.pth、best.pth完全不知道哪个对应哪个配置。我的习惯是在每个训练任务启动时自动创建一个输出目录里面保存一个描述文件import json from datetime import datetime log_info { model: model.__class__.__name__, batch_size: 64, lr: 1e-3, optimizer: Adam, weight_decay: 0.0, start_time: datetime.now().strftime(%Y-%m-%d %H:%M:%S), } with open(run_config.json, w) as f: json.dump(log_info, f, indent2)这样无论过多久你打开任务目录三秒钟之内就能还原出当时的全部配置。配合前面保存的CSV日志和checkpoint目录一套完整的训练记录就齐了。这些事情看似琐碎等到回看某个月前的实验时它们的价值会成倍显现。5.4 我实际一直在用的最小项目结构最后分享一个我自己已经沿用了很长时间的最小训练项目结构。它不复杂但刚好把训练套路固定成可复用的模板project/ config.py # 所有超参数集中定义 dataset.py # 数据加载与预处理 model.py # 模型定义 train.py # 训练主循环包含验证与保存 utils.py # checkpoint保存、日志记录等辅助函数 runs/ # 每次训练的输出目录按时间命名train.py可以翻译成一套完整的流程伪代码所有训练套路的动作都能在里面找到对应的位置读取配置、初始化模型、加载数据、创建优化器与调度器、进入epoch循环、训练、验证、决定是否保存、记录日志文件。这种结构的收益在于换一个新任务时我只需要改config.py和model.py剩下的训练逻辑一行都不用动稳定又安全。零基础读者如果想从“能跑”迈向“用得顺”这个目录结构值得尽早模仿它会帮你把精力集中在真正需要思考的地方也就是网络设计、数据处理和实验结果分析而不是反复重写同一个训练循环。训练套路的本质不是代码本身而是对训练流程的掌控感你知道每一步在干什么知道每一步出了问题该去哪里找答案。把这套流程刻进脑子里之后PyTorch对你来说就不再是文档里零散的API列表而是一套可以随取随用的工作方式。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →