A 60 Minute Blitz 学习笔记:Autograd

本文是 AUTOGRAD: AUTOMATIC DIFFERENTIATION 的学习笔记

autograd 模块是 PyTorch 最核心的模块,提供了对 Tensor 所有运算的自动微分,它是一个“由运行定义”(define-by-run)的框架,也就是说,反向传播是由代码的运行所定义,且在每个迭代中都可以是不同的。

Tensor

torch.Tensorautograd 模块最核心的类。当 tensor 的 .requires_grad 属性被设置为 True 时,所有对该 Tensor 执行的运算都会被记录。当全部的计算执行完毕后,通过调用 .backward() 就可以自动计算所有的梯度。该 tensor 的梯度会被累加到其 .grad 属性中。

如果要停止 tensor 追踪运算记录,可以调用它的 .detach() 方法来把它从运算记录中剔除,同时这个 tensor 未来的运算不再被追踪。

也可以通过 with torch.nograd(): 来停止追踪运算(和占用内存),这在评估模型的时候非常有用,因为模型的可训练参数会将 requires_grad 设置为 True,而在评估模型的时候并不需要使用到梯度。

另一个对自动梯度实现非常重要的类是 Function.

TensorFunction 相互联系并构建起一个无环图,这个无环图会编码完整的运算记录。每个 tensor 都有一个 .grad_fn 属性,该属性是对创建这个 TensorFunction 的引用(用户自己新建的 Tensor.grad_fnNone)。

如果要对输入求梯度,可以调用 Tensor.backward() 方法(对输出张量调用)。如果这个 Tensor 是一个标量(即只含有一个值),在调用 backward() 的时候不需要指明任何参数,否则需要传入一个有相应 shape 的 gradient 参数。

1
import torch

新建一个 tensor x, 设置 requires_grad 为 True

1
2
x = torch.ones(2, 2, requires_grad=True)
print(x)
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)

执行一个张量运算,可以看到 y 的 requires_grad 也是 True

1
2
y = x + 2
print(y)
tensor([[3., 3.],
        [3., 3.]], grad_fn=<AddBackward0>)

y 是由加法运算得到的结果,因此它有 grad.fn 属性

1
print(y.grad_fn)
<AddBackward0 object at 0x7ff5af3443c8>

再对 y 做更多的运算

1
2
3
4
z = y * y * 3
out = z.mean()

print(z, out, sep='n')
tensor([[27., 27.],
        [27., 27.]], grad_fn=<MulBackward0>)
tensor(27., grad_fn=<MeanBackward0>)

.requires_grad_() 方法可以改变 Tensor 的 requires_grad 标志,注意该方法是 In-place 方法,结尾有一个下划线。默认参数是 False

1
2
3
4
5
6
7
8
9
a = torch.randn(2, 2)
a = ((a * 3) / (a - 1))
print(a.requires_grad)

a.requires_grad_(True)
print(a.requires_grad)

b = (a * a).sum()
print(b.grad_fn)
False
True
<SumBackward0 object at 0x7ff563025978>

Gradients

out 张量只包含一个标量值,因此对它调用 out.backward() 相当于调用 out.backward(torch.tensor(1.))

1
2
out.backward() 
print(x.grad) # 查看 x 的梯度
tensor([[4.5000, 4.5000],
        [4.5000, 4.5000]])

推导如下:
$out = frac{1}{4}sum_i z_i$,
$z_i = 3(x_i+2)^2$ ,其中 $z_ibigrrvert_{x_i=1} = 27$.

因此,
$frac{partial o}{partial x_i} = frac{3}{2}(x_i+2)$

所以,
$frac{partial o}{partial x_i}bigrrvert_{x_i=1} = frac{9}{2} = 4.5$.

最终结果是一个所有值均为 4.5 的 2x2 矩阵。

如果是对向量求梯度,则结果将是一个雅可比矩阵。具体说明见 Gradients

1
2
3
4
5
6
7
8
9
x = torch.randn(3, requires_grad=True)

y = x * 2
# 教程原文使用了 y.data.norm(), 但 data 是 0.4 版本的 API,已经不需要使用
while y.norm() < 1000: # y 的范数 > 1000 的时候停止
y = y * 2


print(y)
tensor([ -939.5166, -1193.9800,   809.2726], grad_fn=<MulBackward0>)

y 现在不是一个标量,需要向 backward 方法传入一个向量来求得梯度(关于这个向量的进一步解释,参考PyTorch中的backward):

1
2
3
4
v = torch.tensor([0.1, 1.0, 0.0001], dtype=torch.float)
y.backward(v)

print(x.grad)
tensor([1.0240e+02, 1.0240e+03, 1.0240e-01])

Tensor 包裹在 with torch.no_grad() 上下文中可以使 .requires_grad=True 的张量停止追踪运算记录:

1
2
3
4
5
print(x.requires_grad)
print((x ** 2).requires_grad)

with torch.no_grad():
print((x ** 2).requires_grad)
True
True
False

参考资料: