如何在PyTorch中手动应用梯度？附TensorFlow对应实现对比

问题背景 #

已知在TensorFlow中，可通过以下代码针对网络迭代计算出的loss，对网络net的可训练变量计算梯度并通过优化器应用：

with tf.GradientTape() as tape:              
    grads = tape.gradient(loss, net.trainable_variables)
    optimizer.apply_gradients(zip(grads, net.trainable_variables))

上述方式在使用场景中运行正常。了解PyTorch中的标准梯度应用流程：

optimizer.zero_grad()
loss.backward()
optimizer.step()

但该标准流程不适用于当前场景，请问如何在PyTorch中像TensorFlow那样手动应用梯度？

解决方案 #

在PyTorch里，你可以通过torch.autograd.grad()手动计算梯度，再将梯度传递给优化器完成参数更新，完全对应TensorFlow的手动梯度流程，具体实现如下：

方法一：手动赋值梯度后更新 #

# 清空之前残留的梯度（可选但建议）
optimizer.zero_grad(set_to_none=True)

# 手动计算loss对目标参数的梯度
# create_graph=True用于需要二阶梯度的场景，常规训练设为False即可
grads = torch.autograd.grad(loss, net.trainable_parameters(), create_graph=False)

# 将计算得到的梯度赋值给对应参数的.grad属性
for param, grad in zip(net.trainable_parameters(), grads):
    param.grad = grad

# 应用梯度更新参数
optimizer.step()

方法二：直接传递梯度对给优化器（PyTorch 1.10+支持） #

# 计算loss对目标参数的梯度
grads = torch.autograd.grad(loss, net.trainable_parameters(), create_graph=False)

# 直接将梯度与参数配对后传入step方法，跳过手动赋值.grad的步骤
optimizer.step(lambda: zip(grads, net.trainable_parameters()))

关键细节 #

• 梯度范围控制：如果只需要对部分可训练参数计算梯度，把net.trainable_parameters()替换为你的目标参数列表即可，和TensorFlow指定net.trainable_variables的逻辑一致。
• 梯度清空：如果之前的迭代有残留梯度，务必先调用optimizer.zero_grad()，否则会导致梯度累加，影响参数更新结果。
• 二阶梯度支持：如果需要基于梯度再计算梯度（比如MAML等元学习场景），将torch.autograd.grad()的create_graph参数设为True即可。

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