pytorch中的张量的Autograd.grad（）

Question

我想计算网络中两个张量之间的梯度。输入X张量（批大小xm）通过一组卷积层发送，这些卷积层使我返回并输出Y张量（批大小xn）。

我正在创建一个新的损失，我想知道Y wrt X的梯度。在tensorflow中类似：

tf.gradients（ys = Y，xs = X）

不幸的是，我一直在使用torch.autograd.grad（）进行测试，但是我不知道该怎么做。我收到如下错误：“ RunTimeerror：grad只能为标量输出隐式创建”。

如果我想知道Y wrt X的梯度，torch.autograd.grad（）中的输入应该是什么？

Answer 1

上述解决方案并不完全正确。仅在输出维度为 1 的特殊情况下才正确。

正如文档中提到的， 的输出torch.autograd.grad与导数相关，但实际上并非如此dy/dx。例如，假设您有一个神经网络，输入形状为 的张量(batch_size, input_dim)并输出形状为 的张量(batch_size, output_dim)。输出相对于输入的导数应该具有形状(batch_size, output_dim, input_dim)，但您从中得到的torch.autograd.grad具有形状(batch_size, input_dim)，它是输出维度上的实数导数之和。如果您想要正确的导数，您应该使用torch.autograd.functional.jacobian如下：

import torch
torch.>>> torch.__version__
'1.10.1+cu111'
>>> 

#!/usr/bin/env python
# coding: utf-8

import torch
from torch import nn
import numpy as np


batch_size = 10
hidden_dim = 20
input_dim = 3
output_dim = 2 

model = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.Tanh(), nn.Linear(hidden_dim, output_dim)).double()
x = torch.rand(batch_size, input_dim, requires_grad=True, dtype=torch.float64) #(batch_size, input_dim)
y = model(x) #y: (batch_size, output_dim) 

#using torch.autograd.grad
dydx1 = torch.autograd.grad(y, x, retain_graph=True, grad_outputs=torch.ones_like(y))[0]  #dydx1: (batch_size, input_dim)
print(f' using grad dydx1: {dydx1.shape}')

#using torch.autograd.functional.jacobian
j = torch.autograd.functional.jacobian(lambda t: model(t), x) #j: (batch_size, output_dim, batch_size, input_dim)

#the off-diagonal elements of 0th and 2nd dimension are all zero. So we remove them
dydx2 = torch.diagonal(j, offset=0, dim1=0, dim2=2) #dydx2: (output_dim, input_dim, batch_size)
dydx2 = dydx2.permute(2, 0, 1) #dydx2: (batch_size, output_dim, input_dim)
print(f' using jacobian dydx2: {dydx2.shape}')

#round to 14 decimal digits to avoid noise 
print(np.round((dydx2.sum(dim=1)).numpy(), 14) == np.round(dydx1.numpy(), 14))

输出：

>using grad dydx1: torch.Size([10, 3])

>using jacobian dydx2: torch.Size([10, 2, 3])

#dydx2.sum(dim=1) == dydx1
>[[ True  True  True]
 [ True  True  True]
 [ True  True  True]
 [ True  True  True]
 [ True  True  True]
 [ True  True  True]
 [ True  True  True]
 [ True  True  True]
 [ True  True  True]
 [ True  True  True]]

实际上autograd.grad返回过度输出维度的总和dydx。

如果你真的想使用，torch.autograd.grad有一种低效的方法：

dydx3 = torch.tensor([], dtype=torch.float64)

for i in range(output_dim):
    l = torch.zeros_like(y)
    l[:, i] = 1.
    d = torch.autograd.grad(y, x, retain_graph=True, grad_outputs=l)[0]  #dydx: (batch_size, input_dim)
    dydx3 = torch.concat((dydx3, d.unsqueeze(dim=1)), dim=1)


print(f' dydx3: {dydx3.shape}')
print(np.round(dydx3.numpy(), 14) == np.round(dydx2.numpy(), 14))

输出：

 dydx3: torch.Size([10, 2, 3])
[[[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]

 [[ True  True  True]
  [ True  True  True]]]

我希望它有帮助。

PS 我使用它retain_graph=True是因为多次向后调用。

Answer 2

让我们从简单的工作示例开始，该示例具有简单的损失函数并定期向后移动。我们将建立简短的计算图并对其进行一些梯度计算。

码：

import torch
from torch.autograd import grad
import torch.nn as nn


# Create some dummy data.
x = torch.ones(2, 2, requires_grad=True)
gt = torch.ones_like(x) * 16 - 0.5  # "ground-truths" 

# We will use MSELoss as an example.
loss_fn = nn.MSELoss()

# Do some computations.
v = x + 2
y = v ** 2

# Compute loss.
loss = loss_fn(y, gt)

print(f'Loss: {loss}')

# Now compute gradients:
d_loss_dx = grad(outputs=loss, inputs=x)
print(f'dloss/dx:\n {d_loss_dx}')

输出：

Loss: 42.25
dloss/dx:
(tensor([[-19.5000, -19.5000], [-19.5000, -19.5000]]),)

好的，这可行！现在，让我们尝试重现错误“只能为标量输出隐式创建grad”。如您所见，上一个示例中的损失是一个标量。backward()和grad()通过默认涉及单个标量值：loss.backward(torch.tensor(1.))。如果尝试使用更多值传递张量，则会出现错误。

码：

v = x + 2
y = v ** 2

try:
    dy_hat_dx = grad(outputs=y, inputs=x)
except RuntimeError as err:
    print(err)

输出：

grad can be implicitly created only for scalar outputs

因此，使用grad()时需要指定grad_outputs如下参数：

码：

v = x + 2
y = v ** 2

dy_dx = grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y))
print(f'dy/dx:\n {dy_dx}')

dv_dx = grad(outputs=v, inputs=x, grad_outputs=torch.ones_like(v))
print(f'dv/dx:\n {dv_dx}')

输出：

dy/dx:
(tensor([[6., 6.],[6., 6.]]),)

dv/dx:
(tensor([[1., 1.], [1., 1.]]),)

注意：如果要使用backward()，则只需执行y.backward(torch.ones_like(y))。