一个 PyTorch 模型训练的完整流程

学 PyTorch 时单个概念看懂不难真正容易乱的是完整训练流程。这篇文章先不追求复杂模型只把一条最基础的训练主线串起来。第一步准备数据训练模型前先要把数据整理成模型能吃的形式。通常会经历几步读取原始数据做必要的清洗和预处理转成 Tensor封装成 Dataset用 DataLoader 批量加载如果数据这一步没处理好后面模型再复杂也很难救回来。第二步定义模型PyTorch 里通常会继承nn.Module定义模型from torch import nn ​ class Net(nn.Module): def __init__(self): super().__init__() self.linear nn.Linear(10, 2) ​ def forward(self, x): return self.linear(x)这里__init__定义模型有哪些层forward定义数据怎么流过这些层。第三步选择损失函数损失函数负责衡量模型预测错了多少。分类任务常见loss_fn nn.CrossEntropyLoss()回归任务常见loss_fn nn.MSELoss()损失函数要和任务类型匹配这一点很重要。第四步选择优化器优化器负责根据梯度更新参数。常见写法optimizer torch.optim.Adam(model.parameters(), lr1e-3)这里model.parameters()告诉优化器要更新哪些参数lr是学习率。第五步训练循环最核心的训练循环通常长这样for x, y in train_loader: pred model(x) loss loss_fn(pred, y) ​ optimizer.zero_grad() loss.backward() optimizer.step()这几行非常重要。可以按顺序理解前向传播得到预测计算 loss清空旧梯度反向传播计算新梯度优化器更新参数这就是 PyTorch 训练模型的核心骨架。第六步验证模型训练时还需要在验证集上观察效果。验证阶段通常不需要计算梯度所以会写model.eval() with torch.no_grad(): for x, y in val_loader: pred model(x)这样可以减少显存占用也避免误更新模型。小结一个 PyTorch 训练流程可以压缩成这样数据 - 模型 - loss - backward - optimizer.step - 验证刚开始不要急着堆复杂结构。先把这条主线真正跑通后面再换模型、调参数、加可视化都会轻松很多。技术图把关键链路画清楚可运行实验跑通一个最小训练与验证闭环训练循环真正需要关注的是数据流和状态变化训练阶段计算梯度验证阶段关闭梯度并只统计指标。import torch from torch import nn ​ torch.manual_seed(0) x torch.arange(0, 6, dtypetorch.float32).unsqueeze(1) y 2 * x 1 model nn.Linear(1, 1) optimizer torch.optim.SGD(model.parameters(), lr0.05) loss_fn nn.MSELoss() for epoch in range(101): optimizer.zero_grad() loss loss_fn(model(x), y) loss.backward() optimizer.step() if epoch in (0, 50, 100): print(fepoch{epoch} loss{loss.item():.6f}) with torch.no_grad(): print(fx7 prediction{model(torch.tensor([[7.0]])).item():.3f})运行结果epoch0 loss41.809498 epoch50 loss0.000142 epoch100 loss0.000007 x7 prediction14.996Loss 持续下降模型最终接近真实关系y2x1所以输入 7 时预测接近 15。验证与推理用torch.no_grad()避免无意义的计算图。常见误区验证时只写model.eval()就会关闭梯度。二者职责不同通常还要配合torch.no_grad()。只保存模型对象最方便。更稳妥的做法是保存state_dict、配置和预处理信息。动手练习增加一个验证集每 10 个 epoch 记录训练与验证 loss并画出两条曲线。本文首发于「去你想去的地方」 一个 PyTorch 模型训练的完整流程 | 去你想去的地方完整学习路线、视频版和后续更新请访问原文。