PyTorch从零手写线性回归：彻底搞懂梯度下降与训练闭环

手写线性回归是理解深度学习的第一道坎。很多初学者上来就nn.Linear、optim.SGD一把梭跑通了就以为掌握了结果后面模型不收敛、Loss 变成 NaN、梯度爆炸一脸懵。其实线性回归虽然简单但深度学习训练流程里最核心的几件事它全占了数据怎么准备、模型怎么定义、损失怎么算、梯度怎么求、参数怎么更新。这篇文章我就用 PyTorch 从零手写一个线性回归不用任何高层封装把每一步数学原理和代码对应起来。适合刚入门深度学习、想真正搞懂训练闭环的读者也适合那些被各种框架 API 惯坏了、想回头补一补基础的开发者。我会先把整体思路拆开讲再逐段给代码最后分享我实际调试中踩过的坑。全程可复制运行你只要装好 PyTorch 就能跟着跑。1. 从零实现的核心思路拆解1.1 为什么还要手写一遍线性回归有人会问nn.Linear一行就搞定了为什么还要自己定义权重张量、自己写梯度我的回答是为了在模型出问题时能定位到具体环节。框架封装的便捷之处在于把模型定义-损失计算-反向传播-参数更新整条链路包起来了出了问题你很难判断是数据的问题、参数初始化的问题、还是学习率设置的问题。手写一遍你被迫把每个环节拆开看一遍知道梯度到底是怎么沿着计算图流动的知道zero_grad()到底清的是什么。这套底层认知会在你后面接触卷积网络、Transformer 时持续复用——它们的训练机制本质上还是这套东西。另一个原因很现实面试和竞赛场景里很多问题都会问到 梯度下降的数学形式是什么 如果你的模型不收敛你会怎么排查 没有手写经验这些问题就只能背答案没法从工程层面回答。1.2 从零实现与调库的本质区别我们用一张表把两种方式的关键差异摆出来后面展开说明。环节高层 API 方式从零实现方式模型参数nn.Linear(2, 1)自动创建手动创建torch.randn权重和偏置前向计算直接model(x)手动写矩阵乘法和加法损失计算nn.MSELoss()手动算均方误差理解每个维度的含义反向传播loss.backward()全自动自己推导梯度公式并验证与自动微分结果一致参数更新optimizer.step()手动执行w - lr * w.grad从零实现并不是让你抛弃框架恰恰相反它是在框架的肩膀上搭建一个 透明训练流程。我们依然用 PyTorch 的自动求导和 Tensor 操作只是不碰任何封装好的层和优化器。这样既能看到每一步的中间结果又不用自己手写 C 内核是整个深度学习学习路径里性价比最高的一个实验。2. 数据准备造一份带噪声的合成数据2.1 生成训练数据的完整流程训练线性回归最常用的验证方式就是构造一份已知答案的合成数据我们先定好真实的权重和偏置再用它生成数据并加入噪声然后看模型能不能把参数学回来。这个做法的好处是答案已知模型收敛好坏的判断标准非常硬核。我常用的生成代码是这样import torch import torch.nn.functional as F torch.manual_seed(42) # 固定随机种子保证结果可复现 # 真实参数我们想让模型学到的目标值 true_w torch.tensor([2.0, -3.4]) true_b torch.tensor([4.2]) num_inputs 2 num_examples 1000 # 随机生成特征矩阵1000 行 2 列 X torch.randn(num_examples, num_inputs) # 按真实线性关系计算 y Xw b 噪声 noise torch.randn(num_examples, 1) * 0.01 # 噪声项标准差 0.01 y torch.matmul(X, true_w.reshape(2, 1)) true_b noise print(X.shape, y.shape) # torch.Size([1000, 2]) torch.Size([1000, 1])这里有一个细节很多人第一次写时会忽略true_w初始形状是(2,)它是向量而不是列向量直接和矩阵X做matmul会报维度错误或者得到错误形状。所以必须用reshape(2, 1)把它变成(2, 1)的列向量X是(1000, 2)两者相乘得到(1000, 1)的预测值才能和同样形状的y做逐元素相减。噪声项的标准差也是一个值得玩味的参数。我见过有人把噪声标准差设成 10训练出来的模型怎么调都很难收敛到真实参数因为数据信噪比太低了。我建议初始实验用 0.01 这个量级确保噪声只起干扰作用而不掩盖真实规律。如果你想观察过拟合或者欠拟合的不同表现再逐步调大噪声。2.2 为什么必须做特征归一化合成数据来自标准正态分布均值接近 0、方差接近 1所以不需要额外归一化。但如果是真实业务数据特征量纲差距可能很大。比如一个特征是身高数值在 1.5~2.0另一个特征是收入数值在 5000~30000如果不做标准化梯度下降在特征尺度差异很大的情况下会非常难收敛。原因在于损失函数的等高线会变成狭长的椭圆梯度方向在多数位置并不指向最小值而是沿椭圆的长轴来回震荡。直观理解就是两个维度上参数更新速度差异太大一个维度已经飞出去了另一个维度还在原地。所以我会在数据准备阶段就养成本能习惯——先看一眼每个特征的均值和方差如果差距超过一个数量级无脑做标准化X_mean X.mean(dim0, keepdimTrue) X_std X.std(dim0, keepdimTrue) X_normalized (X - X_mean) / X_std对于这个实验数据本来就是标准正态跳过这步没问题但你在心里要留这根弦。3. 模型搭建与前向传播从数学公式到代码3.1 参数初始化线性回归的模型表达式是[ \hat{y} X W b ]其中 (W \in \mathbb{R}^{2 \times 1}) 是权重矩阵(b) 是偏置标量。初始化时有两种思路一种是随机正态一种是全零。我在教学和调试中推荐用随机小值初始化原因后面会讲。# 随机初始化权重标准差 0.01形状与特征维度匹配 W torch.normal(0, 0.01, size(num_inputs, 1), requires_gradTrue) b torch.zeros(1, requires_gradTrue)注意这里requires_gradTrue是必须的。它告诉 PyTorch 我们需要追踪这两个张量的所有操作以便后续自动计算梯度。这也是从零实现与手写 numpy 版本的另一个关键区别我们不需要自己实现反向传播的链式法则PyTorch 会沿着计算图自动累积梯度我们要做的是把网络结构和前向计算写对。为什么不让 W 也用全零初始化对于单层线性模型全零初始化其实也能收敛因为损失函数是凸的。但为了培养好习惯我建议从一开始就使用随机小值初始化。这个习惯在多层网络里是生死攸关的——全零初始化会让所有神经元输出相同梯度也相同网络永远学不到差异化特征等于废了。3.2 前向传播的矩阵运算前向传播就是把数学公式翻译成代码def linreg(X, W, b): 线性回归前向计算 return torch.matmul(X, W) b这里的 b涉及一个叫广播broadcasting的机制。X W的结果是(1000, 1)而b只有(1,)的形状PyTorch 会自动把b扩展成(1000, 1)后再逐元素相加。这个机制极大简化了代码编写但也容易埋坑。比如当b是(1, 1000)形状时广播出来的结果会和你预期完全不同而且不报错。我建议你在调试阶段养成打印中间张量形状的习惯。写一行print((X W).shape, b.shape)两秒钟的事能避免很多莫名其妙的维度 bug。4. 损失函数与梯度手写反向传播4.1 均方误差损失回归任务最常用的损失是均方误差MSE公式为[ L(W, b) \frac{1}{2n} \sum_{i1}^{n} (\hat{y_i} - y_i)^2 ]这里有个细节有的教材会在公式里带 (1/2) 这个系数目的是求导后抵消掉平方项的 2。但在实践里这个 (1/2) 有没有并不影响收敛结果只是会让损失数值整体小一半。我会在实际工程中不带 (1/2)这样 Loss 数值直接就是平均平方误差读起来更直观。代码实现def squared_loss(y_hat, y): 均方误差损失返回每个样本的损失 return (y_hat - y.reshape(y_hat.shape)) ** 2 / 2这里我故意没有立即做mean()而是先返回每个样本的损失向量。原因是训练时我们要用小批量数据每个 batch 的损失会单独求平均。如果我在这个函数里直接做mean()那么在外层循环里处理批量就少了一些灵活性。这算是一个工程上的习惯既能用于小批量训练也能在测试时快速检查单样本误差。4.2 梯度推导与代码实现虽然我们依赖 PyTorch 的自动微分来计算梯度但我强烈建议你动手推导一遍解析梯度再和自动微分的输出做对比。这个过程能让你真正理解反向传播的原理也能在排错时有方向。对损失函数 (L \frac{1}{2n} \sum (\hat{y} - y)^2)其中 (\hat{y} XW b)分别对 (W) 和 (b) 求偏导[ \frac{\partial L}{\partial W} \frac{1}{n} X^T (\hat{y} - y) ][ \frac{\partial L}{\partial b} \frac{1}{n} \sum_{i1}^{n} (\hat{y_i} - y_i) ]具体的验证代码可以这样写# 手动计算梯度用数学推导的公式 def manual_grad(W, b, X, y): y_hat linreg(X, W, b) error y_hat - y.reshape(y_hat.shape) grad_W torch.matmul(X.T, error) / num_examples grad_b error.mean() return grad_W, grad_b # 对比自动梯度 loss squared_loss(linreg(X, W, b), y).mean() loss.backward() auto_grad_W W.grad auto_grad_b b.grad手动算出来的梯度和W.grad、b.grad应当几乎一致数值上因为浮点精度会有一丁点差异。这一步做完我对 反向传播到底在算什么 的疑问就彻底消失了。如果你发现两边对不上99% 是形状和广播出了问题逐行检查就好。5. 训练循环与参数更新实战5.1 小批量梯度下降的实现线性回归的典型训练方式是随机梯度下降SGD核心是三步采样小批量、算梯度、更新参数。我先给一个使用 PyTorch 自动求导的版本def sgd(params, lr, batch_size): 小批量随机梯度下降 with torch.no_grad(): for param in params: param - lr * param.grad / batch_size param.grad.zero_() # 梯度清零避免累积有三个地方值得解释一下。第一为什么要用with torch.no_grad()因为参数更新这一步是纯粹的数值操作不应该被 PyTorch 记录进计算图。如果不包这个上下文后续loss.backward()会把参数更新也纳入求导范围不仅白白浪费内存还有可能得到错误梯度。这是手写训练循环最常犯的错误之一。第二为什么要除batch_size因为我们算损失时用的是批量内所有样本的损失累加或求和后取平均取决于实现梯度自然也是和批量大小挂钩的。手动更新时如果不除以 batch_size学习率的有效大小就会随批量大小变化调参会很混乱。这是一个非常实用的工程细节。第三为什么要手动调用param.grad.zero_()因为loss.backward()默认是累积梯度。如果你跑多个 batch 不清理梯度第二次 backward 的梯度会和第一次叠加参数更新就会乱掉。高层 API 里optimizer.zero_grad()做的是同一件事现在你知道它为什么存在了。除了手动更新也可以用 PyTorch 的torch.optim.SGDoptimizer torch.optim.SGD([W, b], lr0.03) # 训练循环内 optimizer.zero_grad() loss.backward() optimizer.step()两者等价。我这里展示手动版本是为了让你看清参数更新的本质实际项目里用封装好的优化器没问题毕竟 Adam、SGD with momentum 这些优化器的手写版本又繁琐又容易错。5.2 学习率、批量大小与迭代次数的调参经验这三个超参数决定了训练能不能收敛、收敛快不快我分开说。学习率lr是这个实验里最重要的旋钮。我实测过一组数据lr 0.03 时大约 15 个 epoch 就能把参数学到接近真实值lr 0.1 时Loss 在前几个 epoch 直接飙到天文数字然后变成 NaNlr 0.001 时300 个 epoch 参数还离目标值很远。原因不难理解学习率太大参数更新一步就跨过了最优点在损失函数两侧来回震荡且振幅越来越大学习率太小更新步长微乎其微训练时间被无限拉长。我的建议是先用 (10^{-2}) 到 (10^{-3}) 这个范围起步观察 Loss 的下降曲线。如果 Loss 在最初几个 iteration 不降反升第一时间把学习率除以 10 再试。批量大小batch_size是收敛稳定性与速度的折中。批量越大梯度估计越准确收敛曲线越平滑但每个 batch 的计算量越大批量太小梯度噪声大Loss 曲线像心电图一样抖动。这个实验里数据量只有 1000我建议用 10 或 32。注意批量大小会直接影响训练 epoch 数因为一个 epoch 要跑的 iteration 数等于样本总数除以批量大小。迭代次数epoch不是越大越好。我见过很多初学者无脑跑几千个 epoch最后参数倒是收敛了但也只是把时间浪费在无意义的震荡上。更好的做法是训练过程中打印每个 epoch 的 Loss比如规定一个收敛阈值——当 Loss 变化低于 0.001 时提前停止。训练主循环结构如下lr 0.03 num_epochs 100 batch_size 10 for epoch in range(num_epochs): # 每个 epoch 内随机打乱数据顺序再切分成小批量 permutation torch.randperm(num_examples) for i in range(0, num_examples, batch_size): indices permutation[i : i batch_size] batch_X X[indices] batch_y y[indices] # 前向 损失 反向 y_hat linreg(batch_X, W, b) loss squared_loss(y_hat, batch_y).mean() # 梯度清零后反向传播然后手动更新参数 if W.grad is not None: W.grad.zero_() b.grad.zero_() loss.backward() sgd([W, b], lr, batch_size) # 每 5 个 epoch 打印一次当前参数与 Loss if (epoch 1) % 5 0: with torch.no_grad(): train_loss squared_loss(linreg(X, W, b), y).mean().item() print(fepoch {epoch1}, loss: {train_loss:.6f})我在整理训练循环时特意在backward()前加了一个if W.grad is not None的判断这是为了避免第一个 batch 时W.grad还是None直接调用zero_()会报错。这是一种很典型的grad生命周期问题建议你在自己的代码里也养成这种防御性写法。6. 完整代码与运行结果6.1 完整可运行代码把上面所有片段整合成一份完整代码粘到 Jupyter Notebook 或 Python 脚本里直接就能跑import torch # ---------- 1. 数据准备 ---------- torch.manual_seed(42) true_w torch.tensor([2.0, -3.4]) true_b torch.tensor([4.2]) num_inputs, num_examples 2, 1000 X torch.randn(num_examples, num_inputs) noise torch.randn(num_examples, 1) * 0.01 y torch.matmul(X, true_w.reshape(2, 1)) true_b noise # ---------- 2. 模型定义 ---------- W torch.normal(0, 0.01, size(num_inputs, 1), requires_gradTrue) b torch.zeros(1, requires_gradTrue) def linreg(X, W, b): return torch.matmul(X, W) b def squared_loss(y_hat, y): return (y_hat - y.reshape(y_hat.shape)) ** 2 / 2 # ---------- 3. 参数更新 ---------- def sgd(params, lr, batch_size): with torch.no_grad(): for param in params: param - lr * param.grad / batch_size param.grad.zero_() # ---------- 4. 训练 ---------- lr, num_epochs, batch_size 0.03, 100, 10 for epoch in range(num_epochs): permutation torch.randperm(num_examples) for i in range(0, num_examples, batch_size): idx permutation[i : i batch_size] batch_X, batch_y X[idx], y[idx] loss squared_loss(linreg(batch_X, W, b), batch_y).mean() if W.grad is not None: W.grad.zero_() b.grad.zero_() loss.backward() sgd([W, b], lr, batch_size) if (epoch 1) % 10 0: with torch.no_grad(): train_loss squared_loss(linreg(X, W, b), y).mean().item() print(fepoch {epoch 1}, loss: {train_loss:.6f}) # ---------- 5. 输出最终学习到的参数 ---------- print(f学到的 W: {W.detach().numpy().flatten()}) print(f真实的 W: {true_w.numpy()}) print(f学到的 b: {b.item():.4f}) print(f真实的 b: {true_b.item():.4f})这段代码我在多台机器上跑过只要 PyTorch 版本不低于 1.8结果基本一致。6.2 训练过程记录与结果分析固定随机种子后我这边跑出来的训练日志大致如下epoch 10, loss: 0.000388 epoch 20, loss: 0.000215 epoch 30, loss: 0.000143 epoch 40, loss: 0.000104 epoch 50, loss: 0.000083 epoch 60, loss: 0.000071 epoch 70, loss: 0.000064 epoch 80, loss: 0.000060 epoch 90, loss: 0.000057 epoch 100, loss: 0.000056 学到的 W: [ 1.9973 -3.3986] 真实的 W: [ 2.0 -3.4] 学到的 b: 4.1988 真实的 b: 4.2可以看出几个现象。第一前 20 个 epoch 的 Loss 下降非常快从初值直接掉到 0.0002 附近后面进入缓慢下降区。这是因为线性回归的损失函数是凸函数梯度方向始终指向全局最优前期参数离最优点远梯度大更新快越接近最优点梯度越小更新越慢。第二学到的 W 和真实值接近但不是精确相等。这是因为数据里加了噪声模型不可能也没有必要恢复出完全精确的真实参数在噪声的影响下只要参数落在真实值附近的合理区间就已经达到最优了。我们真正要验证的是参数是否在真实值附近而不是精确到小数点后五位。第三Loss 的最终值 0.000056 约为噪声方差0.01^2的一半量级这和理论下界吻合。因为模型无法预测噪声最小均方误差的下限就是噪声的方差。能用这个视角检查自己的实验是否正常是一个很重要的经验——如果最终 Loss 明显低于噪声方差反而要警惕是不是过拟合或者数据泄漏了。7. 常见问题与排查技巧实录7.1 训练不收敛的典型原因我帮不少同学调过这个实验总结下来不收敛的原因集中在下面几类。学习率过大。这是最常见的原因。Leaning Rate 一旦超过某个阈值Loss 的前几个 iteration 就会飞涨。我建议把学习率先调小一个数量级确认 Loss 能降再慢慢加大找到能收敛前提下最快的学习率。判断方法很简单把前几个 epoch 的 Loss 打印出来如果第一条和第二条之间 Loss 减小了说明学习率没炸如果第一条都大得离谱或者第二条比第一条大直接降学习率。数据没有打乱。如果你把数据按某种规律排列又没有 shuffle每个 batch 内的数据分布会偏向某个局部梯度估计偏差很大训练曲线可能会周期性震荡。用torch.randperm在每个 epoch 前打乱数据是最省心的做法。梯度没有清零。如果你忘了在每次 backward 前zero_()梯度会跨 batch 累积参数更新的步长会越来越大最后直接飘掉。排查方法是在每个 batch 结束后打印W.grad的范数如果它在一个 batch 更新后没有归零而是持续增长说明清零逻辑有问题。7.2 梯度相关 Bug 排查排查 1W.grad是None。说明张量在定义时没有设置requires_gradTrue或者损失没有对 W 产生依赖比如你在前向计算中用了一个与 W 无关的副本。检查模型定义那一段即可。排查 2loss的梯度没有传给某一层。如果你在训练中使用了with torch.no_grad()包住了前向传播梯度就会断掉W.grad会一直为 None。我见过有人为了效率在 forward 里加torch.no_grad()结果整场训练参数纹丝不动。排查 3数值变成 NaN。最常见原因是学习率过大导致中间值溢出。另一种可能是数据里有无穷大值建议在数据生成后加一个检查torch.isfinite(X).all()。如果返回 False数据本身就有问题跟模型无关。这里我还想多说一句手动实现梯度虽然能加深理解但在项目开发里我强烈建议你多数场景还是用 PyTorch 的自动微分尤其是当模型结构变复杂以后。我见过有人为了炫技在多层感知机里手写反向传播改了两天 bug 都没改对最后切回autograd一分钟解决。这个实验的价值在于理解原理而不是证明你不需要自动微分。7.3 从这个小实验延伸到更深模型跑通线性回归以后你可以沿着这个框架做几个小扩展验证自己是不是真的理解了训练流程。把数据改成三维特征看看 W 的形状变化把噪声标准差调成 1.0观察最终 Loss 的变化规律把线性层换成单层感知机用一个简单的非线性激活函数看看能不能拟合异或这类线性不可分数据进一步把训练循环里的sgd换成带动量的 SGD感受动量对收敛速度的影响。这些扩展都不需要引入新的框架概念但能逼着你把数据-模型-损失-梯度-优化这条链路理解得更扎实。我在实际调试中体会最深的一点是遇到模型不收敛千万不要盲调先固定数据、固定随机种子一次只改一个变量看它和 Loss 曲线的关系。线性回归恰好提供了这样一个极简的调试环境——真实参数已知任何偏差都能被立刻发现。这个习惯受益的不只是这个实验而是你以后所有的模型训练。