如何在不影响梯度的前提下修改PyTorch神经网络权重参数

核心原理 #

你只需要在修改权重时用torch.no_grad()上下文管理器关闭梯度追踪，就能保证权重修改不会被记录到计算图中，不影响后续反向传播的梯度计算。修改操作采用原地拷贝copy_实现，不会破坏网络参数的Parameter属性，后续优化器仍可正常更新参数。

完整实现代码 #

保留你原有的网络、优化器、特征、自定义更新函数定义，训练循环修改如下：

for i in range(10):
    opt.zero_grad()
    # 前向传播，用原始权重生成计算图
    out = net(features)
    loss = torch.mean(torch.square(torch.tensor(5) - torch.sum(out)))

    # -------------------- 反向传播前自定义更新权重 --------------------
    with torch.no_grad():
        # 更新第一层fc1的权重和偏置（不需要更新偏置可去掉对应代码）
        net.fc1.weight.copy_(first_layer_update(net.fc1.weight))
        net.fc1.bias.copy_(first_layer_update(net.fc1.bias))
        # 更新第二层fc2的权重和偏置
        net.fc2.weight.copy_(second_layer_update(net.fc2.weight))
        net.fc2.bias.copy_(second_layer_update(net.fc2.bias))
        # 更新第三层fc3的权重和偏置
        net.fc3.weight.copy_(third_layer_update(net.fc3.weight))
        net.fc3.bias.copy_(third_layer_update(net.fc3.bias))
    # -----------------------------------------------------------------

    # 反向传播和优化器步骤不受影响
    loss.backward()
    opt.step()

关键说明 #

• 权重更新的代码必须放在前向传播之后、反向传播之前：前向传播已经用原始权重完成了计算图的构建，此时修改权重不会改变已有的计算图，反向传播计算得到的梯度仍是对应前向传播时权重的梯度，符合你的需求。
• 不要直接用赋值语句（比如net.fc1.weight = first_layer_update(net.fc1.weight)），会把参数的Parameter类型替换为普通张量，导致优化器无法识别可训练参数、梯度计算异常。
• 如果你不需要更新对应层的偏置，直接删掉对应bias的更新代码即可。

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