ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch神经网络训练四步流水线:模型、数据、损失与优化

PyTorch神经网络训练四步流水线:模型、数据、损失与优化 1. 先弄懂这四个步骤是怎么串起来的接触PyTorch也快五年了我一直觉得初学者最容易踩的坑不是某个API记不住而是把整个训练流程当成一盘散沙今天查一下怎么定义网络明天搜一下Dataloader怎么写后天又开始问为什么loss不降。最后代码堆了一大堆却没有一条能跑通。其实当你把“训练一个神经网络”这件事拆开看它本质上就是四个固定动作定义模型、准备数据、选择损失函数和优化器、编写训练循环。这四步不是零散的技巧而是一条流水线每一步的输出正好是下一步的输入。模型告诉你“预测值长什么样”数据给你“正确答案”损失函数衡量“预测和答案差多远”优化器根据这个差距去调整模型参数训练循环则负责把以上所有环节反复执行成千上万次。我写这篇内容的初衷就是想用一个最简单的神经网络例子把这四步从头到尾讲透。顺便回答几个我私下被问过无数遍的问题为什么我的模型训练完损失一直在0.69附近不动为什么同一个模型换个损失函数效果天差地别为什么有人说要先zero_grad()再backward()不管你是刚装了PyTorch准备跑通第一个模型的新手还是写过不少代码但总觉得某些细节“知其然不知其所以然”的进阶选手这篇内容都适合你。我会尽量用大白话把原理讲清楚同时给出可以直接复现的代码。2. 定义模型不只是一个类的封装2.1 nn.Module 和 forward 到底在干什么PyTorch定义模型的标准姿势是继承torch.nn.Module在__init__里声明层结构在forward里定义数据的前向传播路径。很多初学者不理解为什么非得这样写直接在一个普通函数里算不行吗如果你只做一次前向计算当然可以不用类封装。但神经网络的训练反复涉及一个核心操作自动求梯度。PyTorch是在前向传播的过程中动态构建计算图的每当你调用一个张量运算框架就会记录这个运算的输入、输出和梯度函数。nn.Module的核心价值就是帮你把参数注册进一个统一的管理体系里model.parameters()能自动收集所有层的权重model.to(device)能一键把全部参数搬到GPUmodel.train()和model.eval()能统一切换所有层的状态。至于forward它不是被你直接调用的而是通过model(x)触发。这里有个非常关键的细节nn.Module的__call__方法会先执行一些hooks、状态检查然后才真正调用你写的forward。所以千万不要图省事在代码里写model.forward(x)这会绕过PyTorch帮你做的前置处理后面你加个Dropout或者BatchNorm的时候就会发现行为完全不对。2.2 一个能跑通的最小模型长什么样我们用一个最简单的二元分类任务来演示。输入是二维平面上的点如果x0 x1 0就归为正类否则归为负类。这个问题用一条直线就能划分理论上一个线性层就够用但为了展示神经网络的基本形态我还是加一个隐藏层。import torch import torch.nn as nn class SimpleMLP(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(2, 16) self.relu nn.ReLU() self.fc2 nn.Linear(16, 1) def forward(self, x): h self.relu(self.fc1(x)) y self.fc2(h) return y这里nn.Linear(2, 16)的意思是输入维度是2输出维度是16内部维护一个形状为(16, 2)的权重矩阵和一个形状为(16,)的偏置向量。你不需要手动初始化这些参数PyTorch默认会用Kaiming均匀分布初始化线性层这是为了让网络在训练初期前向传播的方差保持稳定。隐藏层之后接了一个ReLU激活函数。至于为什么用ReLU而不是SigmoidReLU计算简单、梯度不会饱和是目前最常用的选择。我知道有些同学喜欢在输出层之后也加一个Sigmoid这里我故意不加原因后面讲损失函数的时候会详细解释。2.3 模型设计的几个实用经验第一能用nn.Sequential就别手写太复杂的forward。比如上面这个模型也可以用self.net nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))来定义。对于线性堆叠结构Sequential的可读性更好但一旦出现跳跃连接、多分支输入就得老老实实写forward了。第二留意__init__里注册子模块的方式。self.fc1 nn.Linear(...)这种直接赋值的方式会被nn.Module自动识别并注册但如果你把层塞进一个普通的Python列表里比如self.layers [nn.Linear(...), nn.ReLU()]这些层就不会被注册parameters()里找不到它们训练时梯度根本更新不到。想用列表就得改用nn.ModuleList。第三别小看model.train()和model.eval()的切换。我见过太多人在推理时忘记切回eval()模式结果模型里有Dropout和BatchNorm的时候预测结果每次都不一样。虽然我们这个小模型没有这两种层但养成好习惯永远不会亏。3. 准备数据别让数据成为模型的绊脚石3.1 从原始数据到PyTorch张量数据准备这件事看似只是“把数组转成Tensor”但实际上它决定了你的模型能不能收敛。我先说最基础的部分PyTorch模型吃的不是NumPy数组、不是Pandas DataFrame而是torch.Tensor。转换很简单torch.from_numpy()或者torch.tensor()都行但有几个细节必须注意。第一是数据类型。模型参数默认是float32你的输入数据最好也是float32。如果你从CSV里读进来的数据是int64或者float64模型会直接报类型不匹配的错。养成习惯训练前统一.float()。第二是数据标准化。我们手搓一个二分类数据集可以直接用torch.randn生成因为标准正态分布的数天生就是“标准化”的。但在真实场景里你要处理的数据量纲乱七八糟有的特征取值范围0到1有的却是几十万。如果直接喂给模型数值大的特征会主导梯度更新模型很容易跑偏。所以无论做什么任务先算均值和标准差然后把数据变成(x - mean) / std的形态这是一个性价比极高的预处理步骤。比如我们这个例子torch.manual_seed(42) X torch.randn(200, 2) y (X[:, 0] X[:, 1] 0).float().unsqueeze(1)torch.manual_seed(42)是为了让结果可复现。y是形状为(200, 1)的列向量注意我用了unsqueeze(1)因为后面的损失函数要求预测值和标签形状一致。这里的形状对齐是初学者高频翻车点多用print(y.shape, logits.shape)核对不会错。3.2 Dataset 与 Dataloader你真的需要它们吗如果你的数据总量就两百条直接全局喂给模型也没问题。但真实项目里数据往往大到内存装不下或者需要按批次打乱顺序训练。这时候就用上了Dataset和Dataloader。Dataset负责定义“怎么取一条数据”Dataloader负责“按什么节奏把数据一批批送进模型”。PyTorch官方推荐的自定义Dataset写法是继承torch.utils.data.Dataset实现__len__和__getitem__两个方法。举个简单例子from torch.utils.data import Dataset, DataLoader, TensorDataset dataset TensorDataset(X, y) loader DataLoader( dataset, batch_size32, shuffleTrue, num_workers2, pin_memoryTrue )TensorDataset是PyTorch提供的最小封装直接传入张量对儿就行了。对于文本、图像这类复杂数据才需要自己写Dataset。3.3 batch_size、shuffle 和 num_workers 的心得batch_size是每次迭代喂给模型的样本数。选太大一次前向反向下来的显存开销高而且梯度方向过于“平均”模型容易陷入尖锐极小值选太小训练速度慢梯度噪声大。经验值一般是32、64、128而且最好满足2的幂次因为很多底层库对这类尺寸有优化。我们这个任务选32就够。shuffle的作用是打乱样本顺序。如果数据集本身有顺序结构比如前一半全是正样本后一半全是负样本不打乱会导致每个batch都偏科模型学出来的决策边界会很抖。训练集开shuffleTrue验证集和测试集开shuffleFalse。num_workers是数据加载的进程数。它不是越大越好我踩过这个坑在小数据集上开16个worker反而因为进程间通信开销把速度拖慢。常规做法是先设成4训练时观察GPU利用率再调整。pin_memoryTrue配合GPU训练能减少数据传输时间但如果数据本身不大这点优化其实可有可无。多说一句如果你的训练数据量很小几百条级别另一个容易忽略的风险是数据划分。我们生成的两百条数据至少要拿出20%做验证集用来观察模型有没有过拟合。交叉验证在小样本场景下更是刚需不要因为懒就省掉。4. 损失函数与优化器最核心的组合拳4.1 分类任务应该选哪个损失函数损失函数的目标只有一个量化“模型预测和真实答案之间的差距”。但“差距”的定义方式有无数种选错了函数模型训练就会南辕北辙。我们这个任务是二分类输出的是“为正类的置信度分数”。这里有两种常见选择第一种是nn.BCELoss。用它的前提是模型输出层必须先经过Sigmoid把分数压缩到0到1之间然后和真实标签计算二值交叉熵。第二种是nn.BCEWithLogitsLoss它内部把Sigmoid和交叉熵融合在了一起你只需要让模型输出原始logit也就是不加Sigmoid的分数。我强烈建议你用第二种。原因有二一是数值稳定性你直接在原始logit上算交叉熵PyTorch内部会做一个log-sum-exp形式的合并计算避免极端值导致梯度爆炸二是方便你少写一层Sigmoid前向传播和反向传播都快一点。多分类任务的对应选择是nn.CrossEntropyLoss它同样内部融合了Softmax你只需要输出每个类别的logit不需要手动加Softmax。很多教程还会提nn.NLLLoss但那是给已经过LogSoftmax的输出用的初学者直接忽略就行。回归任务就完全不同了。预测连续数值一般用nn.MSELoss均方误差或者nn.L1Loss平均绝对误差。MSE对离群点非常敏感L1对离群点更鲁棒。如果拿不准可以先跑一遍几个损失函数看验证集表现。下面这个表是我自己做任务时的常用选择供参考任务类型推荐损失函数输出层操作备注二分类BCEWithLogitsLoss无输出logitlogit0判正类多分类CrossEntropyLoss无输出各类logit取最大logit对应类别多标签BCEWithLogitsLoss无每个标签独立二分类回归MSELoss / L1Loss无输出连续值回归鲁棒SmoothL1Loss无结合MSE和L1优点4.2 优化器从SGD到Adam再到AdamW选好损失函数之后还要选一个优化器来决定“怎么沿着梯度方向更新参数”。最朴素的optim.SGD就是权重减去学习率乘以梯度w w - lr * grad。但SGD收敛慢、容易震荡所以实际中更多人用带Momentum的SGD或者直接用Adam。Adam的核心优势是自适应学习率。它对每个参数单独维护一阶矩和二阶矩估计让梯度大的参数更新得慢一点梯度小的参数更新得快一点。训练初期真的非常省心基本不用手动调学习率。Adam的默认学习率是0.001这个值在大多数任务里都表现不错。但是Adam也有一个隐患它在推断二阶矩的时候用了指数加权平均导致在某些情况下权重衰减的估计有偏。所以后续PyTorch又推出了AdamWAdam with Weight Decay把权重衰减和自适应学习率解耦。PyTorch官方自家的预训练模型和Transformers库默认用的都是AdamW。我自己现在训练小模型也基本只用AdamW。criterion nn.BCEWithLogitsLoss() optimizer torch.optim.AdamW(model.parameters(), lr0.01, weight_decay1e-4)我特意把lr设成0.01来说明对这么简单的任务AdamW用0.01也能很快收敛但如果你要训练的是图像分类那样的大模型学习率一般要降到1e-4到3e-4区间。学习率不是越大越好太大会导致loss震荡甚至直接NaN太小则收敛过慢。我个人的习惯是先按默认值0.001跑10个epoch看趋势loss完全不降就调大到0.01loss震荡就调小到0.0003。4.3 为什么说“损失函数优化器”是一套组合你有没有遇到过这种情况模型结构完全一样但换了个损失函数后训练效果天差地别这是因为损失函数决定了梯度的行为和数值范围而优化器是在这个梯度上做后处理的两者必须匹配。举个直观的例子如果二分类模型输出层加了Sigmoid再用BCELoss那么当预测值接近0或1时Sigmoid的梯度会变得非常平缓也就是所谓的梯度饱和参数更新极其缓慢。而BCEWithLogitsLoss直接在logit上计算能有效规避这个饱和区间。你看同一个优化器搭配不同的损失函数表现的差异就是这么大。还有就是初始化优化器时一定要把model.parameters()传进去。喜欢把参数存成列表再传给优化器的同学要注意如果你重新赋值了模型参数张量比如用了param torch.nn.Parameter(...)用新变量名优化器持有的还是旧张量引用梯度更新就永远不会生效。这是个非常隐蔽的bug我当年排查了整整一个下午。5. 训练循环把前面的零件组装成流水线5.1 一个标准训练循环的解剖终于到了最核心的部分。无论你用什么花哨的训练技巧最终都会落到这个模板上from torch.utils.data import DataLoader, TensorDataset dataset TensorDataset(X, y) train_loader DataLoader(dataset, batch_size32, shuffleTrue) model SimpleMLP() criterion nn.BCEWithLogitsLoss() optimizer torch.optim.AdamW(model.parameters(), lr0.01, weight_decay1e-4) epochs 30 for epoch in range(epochs): model.train() running_loss 0.0 for batch_X, batch_y in train_loader: # 1. 前向传播计算预测 logits model(batch_X) # 2. 计算损失 loss criterion(logits, batch_y) # 3. 梯度清零 optimizer.zero_grad() # 4. 反向传播计算梯度 loss.backward() # 5. 优化器更新参数 optimizer.step() running_loss loss.item() * batch_X.size(0) epoch_loss running_loss / len(train_loader.dataset) print(fEpoch {epoch1:03d}, Loss: {epoch_loss:.4f})这五步就是深度学习训练循环的真身很多人把它简化成“三个大括号”或“三行代码”但我建议你每一步都清楚它的作用。5.2 optimizer.zero_grad() 为什么必须在 backward 之前这是初学者问得最多的问题。PyTorch的Autograd机制是梯度累加的每次backward()计算出来的梯度会累加到参数的.grad属性上。如果你不清零第一次的梯度会一直留着第二次又叠加上去数值越滚越大训练必然崩溃。所以标准顺序是先zero_grad()把旧梯度清零再backward()算出新梯度最后step()更新参数。如果你把zero_grad()放在backward()之后等于刚算出来的梯度又被清零了参数压根不会更新loss会一直纹丝不动。当然这种累加特性也不是纯坑。当你的显存放不下大batch时可以用小batch多跑几次累加梯度模拟一个大batch的效果——但那是进阶玩法初学者先老老实实按标准顺序写。5.3 loss.item() 与标量的故事注意代码里我用的是loss.item()而不是直接打印loss。原因是loss本身是一个带有计算图的张量直接print(loss)会把整个计算图的历史信息都拖着不仅打印出来很啰嗦还会因为保留了loss的引用导致内存越积越多。.item()会把这个张量从计算图中剥离出来返回一个纯Python的float数。这个细节在你写长训练循环时特别重要。如果你把每个epoch的所有loss都保存在一个Python列表里用loss.item()就对了如果直接all_losses.append(loss)列表里每个元素都带着一张计算图跑几轮下来内存就爆了。验证模式也有自己的最佳实践。验证集上不需要计算梯度所以应该包一层torch.no_grad()model.eval() total_correct 0 with torch.no_grad(): for batch_X, batch_y in train_loader: logits model(batch_X) preds (logits 0).float() total_correct (preds batch_y).sum().item() accuracy total_correct / len(train_loader.dataset) print(fTrain accuracy: {accuracy:.2f})torch.no_grad()会关闭梯度计算和计算图构建内存占用大幅下降速度也快不少。同时model.eval()会把BatchNorm和Dropout层的状态切换到推理模式。这两个操作要一起用单用一个都不完整。5.4 给训练循环加一点实用装备跑通了最朴素的循环之后下一步要加点装备不然你盯着控制台里的数字根本不知道模型练得怎么样了。装备一验证集评估。训练集loss下降不代表验证集loss下降过拟合随时会发生。我会在每个epoch结束之后在验证集上算一次loss和准确率只有验证集在变好才算真的在进步。装备二保存最佳模型。不要等到训练完才torch.save而是在验证集loss每创新低的时候把权重存下来。这样即使后面过拟合了你还能把模型恢复到最好的状态。保存方式推荐只存state_dict而不是整个模型因为环境一换整个模型pickle可能就废了。best_loss float(inf) for epoch in range(epochs): # ...训练循环... val_loss evaluate(model, val_loader) if val_loss best_loss: best_loss val_loss torch.save(model.state_dict(), best_model.pth)装备三学习率调整。ReduceLROnPlateau是个免配置的利器当验证集loss连续若干个epoch不下降时它会自动把学习率缩小一半。我常用的参数是patience3, factor0.5效果稳定值得一试。6. 训练中常见的坑我踩过的和你们会踩的6.1 损失不下降卡在某个值不动最常见的一种情况是二分类任务的loss卡在0.693附近。0.693是什么是ln(2)是二分类交叉熵“完全瞎猜”时的理论值。如果你的模型预测每个样本正负的概率都接近0.5loss就是0.693。出现这种现象先别急着调模型结构按顺序排查三个点第一数据标签是不是反了我见过有人把正负样本标签写反模型当然学不到东西。第二数据有没有归一化如果特征范围差异巨大梯度更新会非常不稳定。第三学习率是不是太大或太小学习率太大loss会在一个区间震荡下不去太小则每轮只移动一点点看着就像卡住了。还有一个可能性梯度根本没有传播。检查一下model.parameters()里是否所有参数的grad都不为None如果是None说明你的模型连接在某个地方断了。6.2 loss变成NaN梯度爆炸没跑了NaN是我最不想见的字母组合。出现NaN绝大多数情况是学习率开太大了以致于参数更新一步跨过头导致数值上溢。解决办法很直接把学习率除以10再试。少数情况是你的数据里有NaN值比如读到CSV里某列有缺省模型一顿算就把NaN传染给了所有参数。这种问题要在数据预处理阶段就查干净用torch.isnan(X).any()快速扫描一遍。另外一个诡异来源是你在某些层用了不合适的激活函数比如深层网络里用Sigmoid导致梯度消失或极端饱和从而在反向传播中产生异常数值。这时可以试着换成ReLU族或者加一个BatchNorm层。6.3 显存不足OOM的急救方案显存不够用最直接的办法是把batch_size调小。注意batch_size调小之后你可能需要稍微调大一点学习率来补偿梯度噪声的影响但它们不是严格的线性关系要边测边调。还有一个很隐蔽的显存杀手在训练循环里把loss或中间激活值保存到列表里而忘了.detach()。前面讲过保留张量引用就会保留计算图一个epoch跑下来几十MB累计下去直接爆显存。另外如果你用torch.no_grad()包裹验证代码验证阶段基本不会占显存这条也要检查。6.4 模型在训练集上完美验证集上一塌糊涂这就是典型的过拟合。根本原因通常是模型容量太大或者训练数据太少。解决办法有这么几档先加数据增强或扩充数据次之给模型加DropoutPyTorch里直接nn.Dropout(p0.5)插入到全连接层后面再不然减小模型层数或宽度让模型没那么强的记忆能力最后是在优化器里加weight_decay正则项也就是L2正则化限制权重不要长太大。我在实际项目中还遇到过一种特殊情况训练集和验证集的分布不一致。比如训练数据来自白天验证数据全在晚上即使没有过拟合模型也很难泛化。这种情况下任何正则化手段都救不了得回头审视数据来源。回到最初的问题——四个步骤为什么缺一不可我自己的体会是它们是同一个闭环的四个环节。模型定义决定了假设空间的边界数据决定了空间里有哪些样本损失函数定义了解题的方向优化器和训练循环决定了你沿着这个方向走多快、走多稳。任何一个环节出问题其他环节都会跟着遭殃。所以下次你的模型训练效果不好不要急着换模型结构先按“数据→损失函数→优化器→训练循环”的顺序走一遍体检。这套排查流程我用了五六年救回过无数个半死不活的训练任务。如果你刚开始接触PyTorch也建议先把这套最简单的流程跑熟再去碰分布式训练、混合精度、模型并行那些复杂花样。基础打牢了上层建筑才稳。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进