PyTorch 自动微分 AutoGrad

PyTorch 自动梯度 AutoGrad

1. Python 自动梯度库

Python 中可以做自动梯度（Autograd 或者说自动梯度）的库

• PyTorch

• Tensorflow

• JAX, Github

• Autograd, Github

• Tangent, Github

2 AutoGrad 自动微分

2.1 Auto Gradiant 自动梯度

2.1.1 自动梯度计算过程

自动梯度计算过程：

• (1) 在创建 Tensor 对象时，设定参数requires_grad=True，用于指定微分变量。

  • 或者使用方法 Tensor.requires_grad_(True)

  • 属性 Tensor.requires_grad属性用于判断 Tensor 是否需要计算梯度

• (2) 创建待微分函数的函数关系式

• (3) 利用 Tensor 对象的 .backward() 函数实现自动微分

• (4) 微分结果 Tensor 对象 .grad 属性中

2.1.2 叶子节点的Tensor

1. 只有叶子节点（requires_grad=True 的 Tensor）会计算 Tensor.grad

2. 所有依赖于叶子节点 Tenosr 的 Tensor, 其 requires_grad 属性必定是 True的，但其梯度值只在计算过程中被用到，不会最终存储到 grad 属性中

3. 当使用 backward() 计算梯度之后，计算图（computational graph）会被销毁。如果需要保留计算图，需要设定参数 retain_graph=True

4. 查看梯度值，可以利用 Tensor.register_hook() 打印日志

   Tensor.register_hook(lambda grad: print('grad: ', grad))

2.1.3 Disabling Gradient Tracking

在默认情况下，设置为 requires_grad=True 的 Tensor 会自动追踪计算历史。

用 torch.no_grad() 方法 wrap 相关的代码，会停止最终梯度计算。比如用于模型评估，只需要前向计算。或者使用 Tensor.detach() 方法。

常用情景：

1. frozen some parameters of NNs, for finetuning a pretrained network

2. To speed up computations when only doing forward pass

2.1.4 Accumulative gradients

当先后两次调用 .backward() 时，得到的 Tensor.grad 会不同。PyTorch 默认是累积梯度。

使用 Tensor.grad.zero_() 可以将梯度清零

代码实例：

x = torch.eye(5, requires_grad=True)
y = (x + 1).pow(2)

y.backward(torch.ones_like(x), retain_graph=True)
print("First call\n", x.grad)

y.backward(torch.ones_like(x), retain_graph=True)
print("\n Second call\n", x.grad)

x.grad.zero_()
y.backward(torch.ones_like(x), retain_graph=True)
print("\n Call after zeroing gradients\n", x.grad) 
# equal to the first call

3.2 实例

(1) 标量对标量求导

计算 \(y = f(x) = a x^2 + b x + c\) 在 \(x=0\) 处的梯度，即求：\(f'(x=0) = \left. \dfrac{\mathrm{d} y}{\mathrm{d} x} \right|_{x=0}\) ，其中 \(a = 1, b=-2, c=1\)。

\[f'(x) = \dfrac{\mathrm{d} y}{\mathrm{d} x} = 2ax+b = 2x-2 \quad \rightarrow \quad f'(0) = -2 \]

代码实例：

import numpy as np 
import torch 
x = torch.tensor(0.0, requires_grad=True) # x需要被求导
a, b, c = torch.tensor(1.0),  torch.tensor(-2.0), torch.tensor(1.0)
y = a * torch.pow(x,2) + b * x + c 

y.backward()  
# 等价于   y.backward(gradient=torch.ones_like(x))
# 或等价于 torch.autograd.backward(tensors=y)
print('dy/dx|(x=0)=', x.grad.item())

(2) 矢量对矢量求导

假设自变量为矩阵 \(\mathbf{X}\)（矢量），即 \(\mathbf{Y} = f(\mathbf{X}) = a \mathbf{X}^2 + b \mathbf{X} + c\)

代码实例：

x = torch.tensor([[0.0,0.0], [1.0,2.0]], requires_grad=True) # x 需要被求导
a, b, c = torch.tensor(1.0),  torch.tensor(-2.0), torch.tensor(1.0)
y = a * torch.pow(x,2) + b * x + c 

# 求梯度
y.backward(gradient=torch.ones_like(x))
print(x.grad)

# 前述两行代码 等价于
gradient = torch.tensor([[1.0,1.0], [1.0,1.0]])
z = torch.sum(y*gradient)
z.backward()
print(z.grad)

(3) 求高阶导数

x = torch.tensor(0.0, requires_grad=True) # x 需要被求导
a, b, c = torch.tensor(1.0),  torch.tensor(-2.0), torch.tensor(1.0)
y = a * torch.pow(x,2) + b * x + c
# create_graph 设置为 True 将允许创建更高阶的导数 
dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
print(dy_dx.data)
# 求二阶导数
dy2_dx2 = torch.autograd.grad(dy_dx, x)[0] 
print(dy2_dx2.data)

(4) 对多自变量求导

求 \(\boldsymbol{Y} = \begin{bmatrix} f_1(x_1,x_2) \\ f_2(x_1,x_2) \end{bmatrix} = \begin{bmatrix} x_1x_2 \\ x_1+x_2 \end{bmatrix}\) 对 \(\boldsymbol{X} = \begin{bmatrix} 1 \\ 2 \end{bmatrix}\) 的导数。即，所谓的 Jacobian matrix

\[\frac{\mathrm{d} \boldsymbol{Y}}{\mathrm{d} \boldsymbol{X}} = \begin{bmatrix} \dfrac{\partial y_1}{\partial x_1}, \dfrac{\partial y_1}{\partial x_2} \\ \dfrac{\partial y_2}{\partial x_1}, \dfrac{\partial y_2}{\partial x_2} \end{bmatrix} = \begin{bmatrix} x_2, x_1 \\ 1, 1 \end{bmatrix} = \begin{bmatrix} 2, 1 \\ 1, 1 \end{bmatrix} \]

x1 = torch.tensor(1.0, requires_grad=True) # x需要被求导
x2 = torch.tensor(2.0, requires_grad=True)
y1 = x1 * x2
y2 = x1 + x2
# 允许同时对多个自变量求导数
(dy1_dx1, dy1_dx2) = torch.autograd.grad(outputs=y1, inputs=[x1, x2], retain_graph=True)
print(dy1_dx1, dy1_dx2)
# 如果有多个因变量，相当于把多个因变量的梯度结果求和
(dy12_dx1, dy12_dx2) = torch.autograd.grad(outputs=[y1, y2], inputs=[x1,x2])
print(dy12_dx1, dy12_dx2)

(5) 梯度下降法求最小值

求 \(y = f(x) = a x^2 + b x + c = x^2 - 2x + 1\) 的最小值 \(y^*\) 以及对应的最小值解 \(x^*\)：

\[x \leftarrow x - \eta \cdot \nabla f \]

# 初始化变量（参数），x 为待优化参数，需要被求导
x = torch.tensor(0.0, requires_grad=True)

# 目标函数
def f(x):
    a, b, c = torch.tensor(1.0), torch.tensor(-2.0), torch.tensor(1.0)   # 系数
    y = a * torch.pow(x,2) + b * x + c   # 函数
    return y

lr = 0.1    # 学习率 learning rate
epoch = 20  # 迭代次数

for i in range(epoch):
    x.grad = None  # 梯度清零
    y = f(x)       # 正向输出
    y.backward()   # 计算梯度

    with torch.no_grad():   # 更新参数
        x -= lr * x.grad   # 不等于 x = x - lr * x.grad 
    
    print('Epoch: {0:>3d}, x = {1:8.4f}, y = {2:8.4f}'.format(
          i, x.item(), y.item())
         )

注意：上述代码，有 2 点需要注意

（1）梯度归零。PyTorch 的机制是计算累积梯度，因此，在用梯度对参数进行更新前，需要将参数的梯度归零，即：

x.grad = None
# 或
x.grad.zero_()

（2）在更新自变量 x 时：

• 使用代码 x-=lr*x.grad，不会改变 x 的 requires_grad=True 的属性。

• 如果写作 x=x-lr*x.grad，则会将 x 的 .requires_grad 属性改为 False

或者可以写作：

with torch.no_grad():
    # 正确示例
    x -= lr * x.grad
    print(x.requires_grad)  # True
    
    # 错误示例：
    x = x - lr * x.grad
    print(x.requires_grad)  # False
    
    # 其他方法一：
    x = x - lr * x.grad     # 返回的 Tensor x 的属性 requires_grad 为 False
    x.requires_grad_(True)  # 重新初始化 x, 模型的 x.grad=None
    print(x.requires_grad)  # True
    
    # 其他方法二：
    x.data = x.data - lr * x.grad
    print(x.requires_grad)  # True

• 注意： 初始化（未计算过的）Tensor 的 .grad 默认为 None，使用 x.grad.zero_() 会报错

2.3 使用优化器求最小值

利用 nn.Module 模块建立一个目标函数对象

class ObjFunc(nn.Module):
    def __init__(self, x = None):
        super(ObjFunc, self).__init__()
        # 系数
        self.a = torch.tensor(1.0)
        self.b = torch.tensor(-2.0)
        self.c = torch.tensor(1.0)
        # 初始化 变量（参数）的值
        if x is None:
            x = torch.tensor(0.0)
        self.x = nn.Parameter(x)
    # 正向传播
    def forward(self):
        y = self.a * torch.pow(self.x, 2) + self.b * self.x + self.c 
        return y

lr = 0.1     # 学习率 learning rate
epoch = 30  # 迭代次数

func = ObjFunc()
optimizer = torch.optim.ASGD(func.parameters(), lr=lr)

for i in range(1, epoch+1):
    optimizer.zero_grad()  # 梯度清零
    y = func()             # 正向输出
    y.backward()           # 反向传播，计算梯度
    optimizer.step()       # 更新参数
    print('Epoch: {0:>3d}, x = {1:8.4f}, y = {2:8.4f}'.format(
          i, func.x.item(), y.item())
         )

3. 动态计算图

3.1 PyTorch 中动态计算图简介

Pytorch 的计算图（Computational Graph）是有向无环图（directed acyclic graph，DAG），由节点和边组成，节点表示 torch.Tensor 或者 torch.autograd.Function ，边表示 Tensor 和 Function 之间的依赖关系。

Pytorch 中的计算图是动态图。这里的动态主要有两重含义。

• 第一层含义是：计算图的正向传播是立即执行的。无需等待完整的计算图创建完毕，每条语句都会在计算图中动态添加节点和边，并立即执行正向传播得到计算结果。

• 第二层含义是：计算图在反向传播后立即销毁。下次调用需要重新构建计算图。如果在程序中使用了 backward 方法执行了反向传播，或者利用 torch.autograd.grad 方法计算了梯度，那么创建的计算图会被立即销毁，释放存储空间，下次调用需要重新创建。

3.2 自定义 torch.autograd.Function

参考资料来源于官方 Tutorial 的例子（Ref. [3]）。定义函数 \(f(x) = 0.5\times(5x^3 - 3x)\)，其导数 \(f'(x) = 1.5 (5x^2-3)\)

class Polynomial3(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input): # 前向计算
        ctx.save_for_backward(input)
        return 0.5 * (5 * input ** 3 - 3 * input)
    @staticmethod
    def backward(ctx, grad_output): # 反向传播
        input, = ctx.saved_tensors 
        return grad_output * 1.5 * (5 * input ** 2 - 1)

其中：

1. In the forward pass we receive a Tensor containing the input and return a Tensor containing the output.
2. ctx is a context object that can be used to stash information for backward computation. You can cache arbitrary objects for use in the backward pass using the ctx.save_for_backward method
3. In the backward pass we receive a Tensor containing the gradient of the loss w.r.t. the output, and we need to compute the gradient of the loss w.r.t. the input.

使用自定义的函数

polyfunc = Polynomial3.apply

4 相关函数

Tensor.backward(gradient=None, retain_graph=None, create_graph=False, inputs=None)

参数：

• gradient：Tensor 类型。一般取默认值 None。

• retain_graph ：

• create_graph：bool 类型。True 用于计算高阶导数

• inputs：求导的变量

torch.autograd.backward(tensors, grad_tensors=None, retain_graph=None, create_graph=False, grad_variables=None, inputs=None)

参数：

• tensors：Sequence Tensor or Tensor。需要求导的函数。

torch.autograd.grad()

Tensor.requires_grad_(requires_grad=True)：Params: bool 类型；设置需要计算导数的变量 Tensor

Tensor.requires_grad：Return：bool 类型；判断 Tensor 是否需要计算导数

Tensor.is_leaf：判断 Tensor 是否为计算图的叶子节点

Tensor.grad：Tensor 的微分结果

Tensor.detach()：Return 一个与计算图分离的 Tensor。返回的 Tensor 与原 Tensor 共用内存，用于取消追踪梯度。

torch.no_grad()：取消追踪梯度

Tensor.grad.zero_()：梯度值归零，等价于 Tensor.grad = None

参考文献

Baydin et al., Automatic differentiation in machine learning: a survey, arxiv

Stanford, CS231n Convolutional Neural Networks for Visual Recognition, GitHub

2-2, 自动微分机制, "20天吃掉那只PyTorch", site.

自定义 torch.autograd.Function 实例：Section: Defining new autograd functions, in "Learing PyTorch With Example", Pytorch Tutorial, site

自定义 torch.autograd.Function 实例：二，计算图中的Function, 2-3 动态计算图, in "20 天吃掉那只 PyTorch", site

AUTOMATIC DIFFERENTIATION WITH TORCH.AUTOGRAD, Pytorch Tutorial, site