pytorch 禁止/允许计算局部梯度的操作

在 PyTorch 中，有些操作可以禁止或允许计算局部梯度，这些操作对于梯度计算、优化算法等都有着重要的影响。本文将详细讲解如何禁止/允许计算局部梯度的操作。

禁止计算局部梯度

有些时候，我们不希望某些操作对梯度产生影响，这时候就需要使用 torch.no_grad() 函数来禁止计算局部梯度。示例如下：

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = torch.tensor([4.0, 5.0, 6.0])

# 禁止计算局部梯度
with torch.no_grad():
    z = x + y

# 在禁止计算局部梯度的情况下，修改 x 的值不会影响 z 的结果
x.add_(torch.ones(3))

# 对 z 进行反向传播，不会计算 x 的梯度
z.sum().backward()

print(x.grad)

在这个示例中，我们定义了一个张量 x，通过将 requires_grad 设置为 True，来告诉 PyTorch 需要对该张量进行梯度计算。接着我们定义了一个张量 y，并使用 torch.no_grad() 函数将其与 x 相加赋值给 z，这样在计算 z 的过程中，y 不存在梯度，也就保证了 z 不会对 y 产生梯度。然后，我们修改了 x 的值，但是由于在 torch.no_grad() 的上下文环境中，计算梯度的过程被禁止，所以这时候并不会影响到 z 的结果。最后对 z 进行反向传播，由于 torch.no_grad() 函数的作用，计算梯度的过程会跳过 z 与 y 的计算，只会计算 x 的梯度。

允许计算局部梯度

有些时候，我们希望某些操作对梯度进行影响，这时候就需要使用 torch.enable_grad() 函数来允许计算局部梯度。示例如下：

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# 允许计算局部梯度
with torch.enable_grad():
    y = x * x * x

# 对 y 进行反向传播，会计算 x 的梯度
y.sum().backward()

print(x.grad)

在这个示例中，我们定义了一个张量 x，通过将 requires_grad 设置为 True，来告诉 PyTorch 需要对该张量进行梯度计算。接着，我们使用 torch.enable_grad() 函数来允许计算局部梯度，将 x 的三次方赋值给 y，这样在计算 y 的过程中，x 存在梯度，并会对 y 产生梯度。然后，对 y 进行反向传播，由于 torch.enable_grad() 函数的作用，计算梯度的过程会包括 y 与 x 的计算。

通过这两个示例，我们可以看到，torch.no_grad() 函数和 torch.enable_grad() 函数的作用分别是禁止和允许计算局部梯度，对 PyTorch 中的梯度计算、优化算法等都有着重要的影响。