Pytorch - 获取中间变量/张量的梯度

Question

作为 pytorch 框架 (0.4.1) 中的一个练习，我试图在一个简单的线性层 (Z = X.W + B) 中显示 X（gX 或 dSdX）的梯度。为了简化我的玩具示例，我从 Z 的总和（不是损失）中后退（）。

总而言之，我想要 S=sum(XW+B) 的 gX(dSdX)。

问题是Z（dSdZ）的梯度是None。结果，gX当然也是错的。

import torch
X = torch.tensor([[0.5, 0.3, 2.1], [0.2, 0.1, 1.1]], requires_grad=True)
W = torch.tensor([[2.1, 1.5], [-1.4, 0.5], [0.2, 1.1]])
B = torch.tensor([1.1, -0.3])
Z = torch.nn.functional.linear(X, weight=W.t(), bias=B)
S = torch.sum(Z)
S.backward()
print("Z:\n", Z)
print("gZ:\n", Z.grad)
print("gX:\n", X.grad)

结果：

Z:
 tensor([[2.1500, 2.9100],
        [1.6000, 1.2600]], grad_fn=<ThAddmmBackward>)
gZ:
 None
gX:
 tensor([[ 3.6000, -0.9000,  1.3000],
        [ 3.6000, -0.9000,  1.3000]])

如果我使用 nn.Module，我得到完全相同的结果，如下所示：

class Net1Linear(torch.nn.Module):
    def __init__(self, wi, wo,W,B):
        super(Net1Linear, self).__init__()
        self.linear1 = torch.nn.Linear(wi, wo)
        self.linear1.weight = torch.nn.Parameter(W.t())
        self.linear1.bias = torch.nn.Parameter(B)
    def forward(self, x):
        return self.linear1(x)
net = Net1Linear(3,2,W,B)
Z = net(X)
S = torch.sum(Z)
S.backward()
print("Z:\n", Z)
print("gZ:\n", Z.grad)
print("gX:\n", X.grad)

Answer 1

首先，您只计算通过将requires_grad 设置为True 来启用梯度的张量的梯度。

所以你的输出和预期的一样。你得到X的渐变。

出于性能原因，PyTorch 不会保存中间结果的梯度。因此，您将只获得您将 requires_grad 设置为 True 的那些张量的梯度。

但是，您可以在计算过程中使用register_hook 提取中间梯度或手动保存。这里我只是保存到张量Z的grad变量中：

import torch

# function to extract grad
def set_grad(var):
    def hook(grad):
        var.grad = grad
    return hook

X = torch.tensor([[0.5, 0.3, 2.1], [0.2, 0.1, 1.1]], requires_grad=True)
W = torch.tensor([[2.1, 1.5], [-1.4, 0.5], [0.2, 1.1]])
B = torch.tensor([1.1, -0.3])
Z = torch.nn.functional.linear(X, weight=W.t(), bias=B)

# register_hook for Z
Z.register_hook(set_grad(Z))

S = torch.sum(Z)
S.backward()
print("Z:\n", Z)
print("gZ:\n", Z.grad)
print("gX:\n", X.grad)

这将输出：

Z:
 tensor([[2.1500, 2.9100],
        [1.6000, 1.2600]], grad_fn=<ThAddmmBackward>)
gZ:
 tensor([[1., 1.],
        [1., 1.]])
gX:
 tensor([[ 3.6000, -0.9000,  1.3000],
        [ 3.6000, -0.9000,  1.3000]])

希望这会有所帮助！

顺便说一句：通常你会希望为你的参数激活梯度 - 所以你的权重和偏差。因为在使用优化器时你现在要做的是改变你的输入X，而不是你的权重W 和偏差B。因此，在这种情况下，W 和 B 通常会激活渐变。

Answer 2

blue-phoenox，感谢您的回答。我很高兴听说 register_hook()。

让我认为我有一个错误的 gX 是因为它与 X 的值无关。我必须做数学才能理解它。但是使用 CCE 损失而不是 SUM 会使事情变得更加清晰。所以我为那些可能感兴趣的人更新了这个例子。在这种情况下，使用 SUM 是个坏主意。

T_dec = torch.tensor([0, 1])
X = torch.tensor([[0.5, 0.8, 2.1], [0.7, 0.1, 1.1]], requires_grad=True)
W = torch.tensor([[2.7, 0.5], [-1.4, 0.5], [0.2, 1.1]])
B = torch.tensor([1.1, -0.3])
Z = torch.nn.functional.linear(X, weight=W.t(), bias=B)
print("Z:\n", Z)
L = torch.nn.CrossEntropyLoss()(Z,T_dec)
Z.register_hook(lambda gZ: print("gZ:\n",gZ))
L.backward()
print("gX:\n", X.grad)

结果：

Z:
 tensor([[1.7500, 2.6600],
        [3.0700, 1.3100]], grad_fn=<ThAddmmBackward>)
gZ:
 tensor([[-0.3565,  0.3565],
        [ 0.4266, -0.4266]])
gX:
 tensor([[-0.7843,  0.6774,  0.3209],
        [ 0.9385, -0.8105, -0.3839]])

Answer 3

有一个更简单的方法。只需使用retain_grad():

https://pytorch.org/docs/stable/autograd.html#torch.Tensor.retain_grad

Z.retain_grad()

在致电backward()之前