Pytorch autograd.backward理解

通常在训练的时候, 最后的loss是一个标量, 无脑使用loss.backward()进行反向传播计算梯度即可. 但是碰到有些代码中出现了多个loss, 比如这样的代码:

torch.autograd.backward([loss0, loss1], [torch.ones_like(loss0), torch.ones_like(loss1)])

当loss是一个标量值的时候

此时反向传播求梯度比较好理解, 假设

>>> w = torch.tensor([[1.0, 2.0]], requires_grad=True)	# 可以看做模型中的某个参数 dim=1x2
>>> x  = torch.tensor([[1.0], [1.0]])					# 可以看做模型的训练数据 dim=2x1
>>> loss = torch.mm(w, x)															
>>> loss.backward()								# 也可以执行torch.autograd.backward(loss), 同样的效果
tensor([[3.]], grad_fn=<MmBackward>)
>>> w.grad										# 查看 w的梯度
tensor([1.0, 1.0])				

$$\begin{gathered} loss=w_0\ast x_0+w_1\ast x_1 \\ \frac{\partial loss}{\partial w}=[x_0,x_1] \end{gathered}$$

当loss是一个向量的时候

比如

>>> w = torch.tensor([[1.0, 2.0]], requires_grad=True)			# 可以看做模型中的某个参数 dim=1x2
>>> x  = torch.tensor([[1.0, 1.0], [2.0, 2.0]])							# 可以看做模型的训练数据 dim=2x2
>>> loss = torch.mm(w, x)			
>>> loss
tensor([[5., 5.]], grad_fn=<MmBackward>)																				
>>> loss.backward()								# 此时会报错
...
RuntimeError: grad can be implicitly created only for scalar outputs

报错的意思是 只有标量才能隐式的计算梯度(言下之意是向量需要显示的计算); 那么所谓"显示", 应该就是本文开头 torch.autograd.backward( , )中第二个参数的作用了吧.

# 在本例中呢, 就是以下代码
loss.backward(torch.ones_like(loss))
# 同torch.autograd.backward(loss, torch.ones_like(loss))

• 求和变成标量再求梯度? 先不管这个"显示参数", 如果我把两个loss加起来, 变成一个标量, 应该就可以了隐式的计算梯度了吧.

>>> loss_sum = loss.sum()
>>> loss_sum
tensor(10., grad_fn=<SumBackward0>)
>>> loss_sum.backward(retain_graph=True)			# 设置retain_graph=True, 方便后续实验
>>> w.grad
tensor([[2., 4.]])

• 给定grad_tensor? 继续实验, 如果给定"显示参数"呢? 就是查看torch.autograd.backward说明里面的第二个参数: “grad_tensor”. 可以看出和 "求和变成标量再求梯度"效果是一样的.

>>> w.grad *= 0			#先 将梯度清零
>>> loss.backward(torch.ones_like(loss))
>>> w.grad
tensor([[2., 4.]])

• 分别对每个loss单独求梯度? 既然loss是一个向量, 那么我向量中的每一个元素都是标量, 我可以分别求梯度.

>>> w.grad *= 0			# 先 将梯度清零
>>> loss0, loss1 = loss[0]			
>>> loss0.backward(retain_graph=True)		# 此时必须retain_graph=True, 不然无法继续叠加
>>> loss1.backward(retain_graph=True)
tensor([[2., 4.]])

综上可以看出, 对于向量或者矩阵执行backward操作的时候, 类似与将其进行了"加权求和", 然后对结果进行反向传播. 例子中指定的grad_tensor的元素值全为1, 也可为0或者其他值.

本文参考了Pytorch autograd,backward详解