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

Question

作为 pytorch 框架（0.4.1）中的练习，我尝试在简单的线性层（Z = XW + 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_gradto启用梯度True。

所以你的输出正如人们所期望的那样。您将得到 的梯度X。

出于性能原因，PyTorch 不保存中间结果的梯度。requires_grad所以你只会得到你设置的那些张量的梯度True。

但是，您可以register_hook在计算过程中提取中间梯度或手动保存它。这里我只是将其保存到gradtensor变量中Z：

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

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

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

Z.retain_grad()

打电话之前backward()