高效理解pytorch的backward需要scalar outputs

利用backward时 ， 可能经常遇到错误
RuntimeError: grad can be implicitly created only for scalar outputs

理解的最好方式就是通过具体的例子和代码，下面我将从几个例子和对应的推导来理解backward机制。

backward with scalar output

import  torch
from    matplotlib import pyplot as plt
import  torch.nn.functional as F
# simple gradient
# tensor传数据 Tensor传维度 只有浮点型才能计算梯度
a = torch.tensor([2, 3],dtype=torch.float64, requires_grad=True)
b = a + 3
c = b * b * 3
out = c.mean()
out.backward()
print('*'*10)
print('=====simple gradient======')
print('input')
print(a.data)
print('compute result is')
print(out)
print('input gradients are')
print(a.grad)

结果如下：

**********
=====simple gradient======
input
tensor([2., 3.], dtype=torch.float64)
compute result is
tensor(91.5000, dtype=torch.float64, grad_fn=<MeanBackward0>)
input gradients are
tensor([15., 18.], dtype=torch.float64)

上面做的事情如下：

需要注意的是：out是一个scalar值，也就是一维的数据[1]，而backward是针对scalar输出，也就是说如果数据是[3]就会报错！

backward on non-scalar output

那么如果出现了输出结果不是一维的，但是我还需要求梯度呢？

n.backward ( p )

# backward on non-scalar output
m = torch.tensor([2, 4],dtype=torch.float64, requires_grad=True)
n = torch.zeros(1,2)
n[0,0] = m[0] ** 2
n[0,1] = m[1] ** 3
# n.shape=[1,2]  m.shape=[2] m.unsqueeze(0).shape=[1,2]
# backward只对标量tensor[1]进行求导
n.backward(m.unsqueeze(0))
print('*'*10)
print('=====non scalar output======')
print('input')
print(m)
print('input gradients are')
print(m.grad)

结果

**********
=====non scalar output======
input
tensor([2., 4.], dtype=torch.float64, requires_grad=True)
input gradients are
tensor([  8., 192.], dtype=torch.float64)

n.backward (p) 相等于对n中各位数字求导数后，然后乘以p，作为系数。以本例为例，相当于下图中的右侧。

n.sum().backward()

# backward on non-scalar output
m = torch.tensor([2, 4],dtype=torch.float64, requires_grad=True)
n = torch.zeros(1,2)
n[0,0] = m[0] ** 2
n[0,1] = m[1] ** 3
# n.shape=[1,2] , p.shape=[1,2]
# backward只对标量tensor[1]进行求导 直接在backward传入值就是乘法关系 n * p
n.sum().backward()
print('*'*10)
print('=====non scalar output======')
print('input')
print(m)
print('input gradients are')
print(m.grad)

结果

**********
=====non scalar output======
input
tensor([2., 4.], dtype=torch.float64, requires_grad=True)
input gradients are
tensor([ 4., 48.], dtype=torch.float64)

这里的sum，就是当n的各项导数相加。其实相当于p=[1,1]，即等效代码如下：

# backward on non-scalar output
m = torch.tensor([2, 4],dtype=torch.float64, requires_grad=True)
n = torch.zeros(1,2)
n[0,0] = m[0] ** 2
n[0,1] = m[1] ** 3
p = torch.tensor([1.,1.]).unsqueeze(0)
# n.shape=[1,2] , p.shape=[1,2]
# backward只对标量tensor[1]进行求导 直接在backward传入值就是乘法关系 
n.backward(p)
print('*'*10)
print('=====non scalar output======')
print('input')
print(m)
print('input gradients are')
print(m.grad)

n.mean().backward()

# backward on non-scalar output
m = torch.tensor([2, 4],dtype=torch.float64, requires_grad=True)
n = torch.zeros(1,2)
n[0,0] = m[0] ** 2
n[0,1] = m[1] ** 3
# n.shape=[1,2] , p.shape=[1,2]
# backward只对标量tensor[1]进行求导 直接在backward传入值就是乘法关系 n * p
n.mean().backward()
print('*'*10)
print('=====non scalar output======')
print('input')
print(m)
print('input gradients are')
print(m.grad)

结果

**********
=====non scalar output======
input
tensor([2., 4.], dtype=torch.float64, requires_grad=True)
input gradients are
tensor([ 2., 24.], dtype=torch.float64)

jacobian矩阵

j = torch.zeros(2 ,2)
k = torch.zeros(1, 2)
m = torch.tensor([2, 4],dtype=torch.float64, requires_grad=True)
k[0, 0] = m[0] ** 2 + 3 * m[1]
k[0, 1] = m[1] ** 2 + 2 * m[0]
k.backward(torch.FloatTensor([[1, 0]]),retain_graph=True)
j[0, :] = m.grad
m.grad.zero_() # 梯度清零
k.backward(torch.FloatTensor([[0, 1]]))
j[1, :] = m.grad
print('jacobian matrix is')
print(j)

结果

jacobian matrix is
tensor([[4., 3.],
        [2., 8.]])

上面的代码可以用下图表示清楚
需要注意雅克比矩阵需要分开求，如果k.backward([1,1])则结果为[6,11]。需要注意。