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

Pytorch - Getting gradient for intermediate variables / tensors

作为 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)

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

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

的渐变

出于性能原因,PyTorch 不会保存中间结果的梯度。所以你只会得到你设置 requires_gradTrue.

的那些张量的梯度

但是您可以使用 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。所以在这种情况下,通常会为 WB 激活梯度。

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

让我认为我有一个错误的 gX 的原因是它独立于 X 的值。我将不得不做数学来理解它。但是使用 CCE Loss 而不是 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]])

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

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

Z.retain_grad()

在调用 backward()

之前