PyTorch中冻结中间层：torch.no_grad()下反向传播是否影响第一层？

我设计了一个简单的神经网络，结构如下：

lin0 = nn.Linear(2, 2)
lin1 = nn.Linear(2, 2)
lin2 = nn.Linear(2, 2)

我的目标是冻结第二层（lin1），同时让第一层和第三层的权重可更新。我尝试在forward函数中编写如下代码：

x = lin0(x)
with torch.no_grad():    
    x = lin1(x)
x = lin2(x)

这应该能冻结lin1的所有参数，但我想知道反向传播是否仍会作用于第一层？如果会，权重是如何更新的？

反向传播会作用于第一层lin0，权重更新逻辑和正常训练时一致，具体过程如下：

• 前向传播阶段，lin0的计算未被torch.no_grad()包裹，会保留梯度计算的相关信息；而lin1的计算被这个上下文管理器屏蔽，不会生成梯度图，因此lin1的参数完全不会参与后续的梯度更新。
• 反向传播时，损失的梯度会从lin2开始反向传导，经过lin1时因为没有梯度信息直接跳过，继续传递到lin0的输出节点，接着计算lin0权重和偏置对应的梯度值，之后优化器就会用这些梯度来更新lin0的参数。
• 简单来说，lin1在这里就是一个固定的“特征转换器”，它的输出只是对lin0的结果做了一次无梯度的变换，完全不干扰lin0和lin2的梯度传导与参数更新流程。

另外提个更规范的冻结参数方法：直接将lin1的所有参数的requires_grad设为False，代码如下：

for param in lin1.parameters():
    param.requires_grad = False

这种方式比在forward里套torch.no_grad()更稳妥，不管前向传播的代码怎么写，lin1的参数都不会参与梯度计算。

内容的提问来源于stack exchange，提问作者earnest