最新公告
  • 欢迎光临数据科学与编程,我们是数据学科学兴趣交流小组立即加入我们
  • PyTorch 源码解读之 torch.autograd

    前言

    本篇笔记以介绍 pytorch 中的 autograd 模块功能为主,主要涉及 torch/autograd 下代码,不涉及底层的 C++ 实现。本文涉及的源码以 PyTorch 1.7 为准。

    • torch.autograd.function (函数的反向传播)
    • torch.autograd.functional (计算图的反向传播)
    • torch.autograd.gradcheck (数值梯度检查)
    • torch.autograd.anomaly_mode (在自动求导时检测错误产生路径)
    • torch.autograd.grad_mode (设置是否需要梯度)
    • model.eval() 与 torch.no_grad()
    • torch.autograd.profiler (提供 function 级别的统计信息)

    torch.autograd.function (函数的反向传播)

    我们在构建网络的时候,通常使用 pytorch 所提供的nn.Module (例如nn.Conv2dnn.ReLU等)作为基本单元。而这些 Module 通常是包裹 autograd function,以其作为真正实现的部分。例如nn.ReLU 实际使用torch.nn.functional.reluF.relu):

    class ReLU(Module):
    __constants__ = [‘inplace’]
    inplace: bool

    def __init__(self, inplace: bool = False):
    super(ReLU, self).__init__()
    self.inplace = inplace

    def forward(self, input: Tensor) -> Tensor:
    return F.relu(input, inplace=self.inplace)

    这里的F.relu类型为function,若再剥开一层,其实际包裹的函数类型为builtin_function_or_method,这也是真正完成运算的部分。这些部分通常使用 C++ 实现(如ATen)。至此我们知道,一个模型的运算部分由 autograd functions 组成,这些 autograd functions 内部定义了 forward,backward 用以描述前向和梯度反传的过程,组合后可以实现整个模型的前向和梯度反传。以torch.autograd.function中所定义的Function类为基类,我们可以实现自定义的autograd function,所实现的 function 需包含forwardbackward两个方法。以下以ExpGradCoeff两个自定义 autograd function 为例进行讲解:

    @staticmethod
    def forward(ctx, i): # 模型前向
    result = i.exp()
    ctx.save_for_backward(result) # 保存所需内容,以备backward时使用,所需的结果会被保存在saved_tensors元组中;此处仅能保存tensor类型变量,若其余类型变量(Int等),可直接赋予ctx作为成员变量,也可以达到保存效果
    return result

    @staticmethod
    def backward(ctx, grad_output): # 模型梯度反传
    result, = ctx.saved_tensors # 取出forward中保存的result
    return grad_output * result # 计算梯度并返回

    # 尝试使用
    x = torch.tensor([1.], requires_grad=True) # 需要设置tensor的requires_grad属性为True,才会进行梯度反传
    ret = Exp.apply(x) # 使用apply方法调用自定义autograd function
    print(ret) # tensor([2.7183], grad_fn=<ExpBackward>)
    ret.backward() # 反传梯度
    print(x.grad) # tensor([2.7183])

    Exp 函数的前向很简单,直接调用 tensor 的成员方法exp即可。反向时,我们知道  , 因此我们直接使用  乘以grad_output即得梯度。我们发现,我们自定义的函数Exp正确地进行了前向与反向。同时我们还注意到,前向后所得的结果包含了grad_fn属性,这一属性指向用于计算其梯度的函数(即Expbackward函数)。关于这点,在接下来的部分会有更详细的说明。接下来我们看另一个函数GradCoeff,其功能是反传梯度时乘以一个自定义系数。

    @staticmethod
    def forward(ctx, x, coeff): # 模型前向
    ctx.coeff = coeff # 将coeff存为ctx的成员变量
    return x.view_as(x)

    @staticmethod
    def backward(ctx, grad_output): # 模型梯度反传
    return ctx.coeff * grad_output, None # backward的输出个数,应与forward的输入个数相同,此处coeff不需要梯度,因此返回None

    # 尝试使用
    x = torch.tensor([2.], requires_grad=True)
    ret = GradCoeff.apply(x, 0.1) # 前向需要同时提供x及coeff,设置coeff为-0.1
    ret = ret ** 2
    print(ret) # tensor([4.], grad_fn=<PowBackward0>)
    ret.backward()
    print(x.grad) # tensor([-0.4000]),梯度已乘以相应系数

    torch.autograd.functional (计算图的反向传播)

    在此前一节,我们描述了单个函数的反向传播,以及如何编写定制的 autograd function。在这一节中,我们简单介绍 pytorch 中所提供的计算图反向传播的接口。

    在训练过程中,我们通常利用 prediction 和 groundtruth label 来计算 loss(loss 的类型为Tensor),随后调用loss.backward()进行梯度反传。而 Tensor 类的backward方法,实际调用的就是torch.autograd.backward这一接口。这一 python 接口实现了计算图级的反向传播。

    def backward(self, gradient=None, retain_graph=None, create_graph=False):
    relevant_args = (self,)

    torch.autograd.backward(self, gradient, retain_graph, create_graph)
    # gradient: 形状与tensor一致,可以理解为链式求导的中间结果,若tensor标量,可以省略(默认为1)
    # retain_graph: 多次反向传播时梯度累加。反向传播的中间缓存会被清空,为进行多次反向传播需指定retain_graph=True来保存这些缓存。
    # create_graph: 为反向传播的过程同样建立计算图,可用于计算二阶导

    在 pytorch 实现中,autograd 会随着用户的操作,记录生成当前 variable 的所有操作,并建立一个有向无环图 (DAG)。图中记录了操作Function,每一个变量在图中的位置可通过其grad_fn属性在图中的位置推测得到。在反向传播过程中,autograd 沿着这个图从当前变量(根节点 F)溯源,可以利用链式求导法则计算所有叶子节点的梯度。每一个前向传播操作的函数都有与之对应的反向传播函数用来计算输入的各个 variable 的梯度,这些函数的函数名通常以Backward结尾。我们构建一个简化的计算图,并以此为例进行简单介绍。

    我们再来看下面的计算图,并在这个计算图上模拟 autograd 所做的工作:

    F.manual_grad = torch.tensor(1) # 我们用manual_grad表示,在已知计算图结构的情况下,我们模拟autograd过程手动算得的梯度
    D.manual_grad, E.manual_grad = F.grad_fn(F.manual_grad)
    C.manual_grad, tmp2 = E.grad_fn(E.manual_grad)
    D.manual_grad = D.manual_grad + tmp2 # 这里我们先完成D上的梯度累加,再进行反传
    A.manual_grad = C.grad_fn(C.manual_grad)
    B.manual_grad = D.grad_fn(D.manual_grad) # (tensor([24.], grad_fn=<MulBackward0>), tensor([40.], grad_fn=<MulBackward0>))

    下面,我们编写一个简单的函数,在这个计算图上进行autograd,并验证结果是否正确:

    A = torch.tensor([3.], requires_grad=True)
    B = torch.tensor([2.], requires_grad=True)
    C = A ** 2
    D = B ** 2
    E = C * D
    F = D + E

    autograd(F.grad_fn, torch.tensor(1))
    print(A.auto_grad, B.auto_grad) # tensor(24., grad_fn=<UnbindBackward>) tensor(40., grad_fn=<AddBackward0>)

    # 这一autograd同样可作用于编写的模型,我们将会看到,它与pytorch自带的backward产生了同样的结果
    from torch import nn

    class MLP(nn.Module):
    def __init__(self):
    super().__init__()
    self.fc1 = nn.Linear(10, 5)
    self.relu = nn.ReLU()
    self.fc2 = nn.Linear(5, 2)
    self.fc3 = nn.Linear(5, 2)
    self.fc4 = nn.Linear(2, 2)

    def forward(self, x):
    x = self.fc1(x)
    x = self.relu(x)
    x1 = self.fc2(x)
    x2 = self.fc3(x)
    x2 = self.relu(x2)
    x2 = self.fc4(x2)
    return x1 + x2

    x = torch.ones([10], requires_grad=True)
    mlp = MLP()
    mlp_state_dict = mlp.state_dict()

    # 自定义autograd
    mlp = MLP()
    mlp.load_state_dict(mlp_state_dict)
    y = mlp(x)
    z = torch.sum(y)
    autograd(z.grad_fn, torch.tensor(1.))
    print(x.auto_grad) # tensor([-0.0121, 0.0055, -0.0756, -0.0747, 0.0134, 0.0867, -0.0546, 0.1121, -0.0934, -0.1046], grad_fn=<AddBackward0>)

    mlp = MLP()
    mlp.load_state_dict(mlp_state_dict)
    y = mlp(x)
    z = torch.sum(y)
    z.backward()
    print(x.grad) # tensor([-0.0121, 0.0055, -0.0756, -0.0747, 0.0134, 0.0867, -0.0546, 0.1121, -0.0934, -0.1046])

    pytorch 使用动态图,它的计算图在每次前向传播时都是从头开始构建,所以它能够使用python 控制语句(如 for、if 等)根据需求创建计算图。下面提供一个例子:

    x = torch.tensor([0.3071, 1.1043, 1.3605, 0.3471], requires_grad=True)
    y = f(x) # y = x[0]*x[1]*x[2]
    y.backward()
    print(x.grad) # tensor([1.5023, 0.4178, 0.3391, 0.0000])

    x = torch.tensor([ 1.2817, 1.7840, 1.7033, 0.1302], requires_grad=True)
    y = f(x) # y = x[0]*x[1]*x[3]
    y.backward()
    print(x.grad) # tensor([0.2323, 0.1669, 0.0000, 2.2866])

    此前的例子使用的是Tensor.backward()接口(内部调用autograd.backward),下面我们来介绍autograd提供的jacobian()hessian()接口,并直接利用其进行自动微分。这两个函数的输入为运算函数(接受输入 tensor,返回输出 tensor)和输入 tensor,返回 jacobian 和 hessian 矩阵。对于jacobian接口,输入输出均可以为 n 维张量,对于hessian接口,输出必需为一标量。jacobian返回的张量 shape 为output_dim x input_dim(若函数输出为标量,则 output_dim 可省略),hessian返回的张量为input_dim x input_dim。除此之外,这两个自动微分接口同时支持运算函数接收和输出多个 tensor。

    fc = Linear(4, 2)
    pool = AvgPool2d(kernel_size=2)

    def scalar_func(x):
    y = x ** 2
    z = torch.sum(y)
    return z

    def vector_func(x):
    y = fc(x)
    return y

    def mat_func(x):
    x = x.reshape((1, 1,) + x.shape)
    x = pool(x)
    x = x.reshape(x.shape[2:])
    return x ** 2

    vector_input = torch.randn(4, requires_grad=True)
    mat_input = torch.randn((4, 4), requires_grad=True)

    j = jacobian(scalar_func, vector_input)
    assert j.shape == (4, )
    assert torch.all(jacobian(scalar_func, vector_input) == 2 * vector_input)
    h = hessian(scalar_func, vector_input)
    assert h.shape == (4, 4)
    assert torch.all(hessian(scalar_func, vector_input) == 2 * torch.eye(4))
    j = jacobian(vector_func, vector_input)
    assert j.shape == (2, 4)
    assert torch.all(j == fc.weight)
    j = jacobian(mat_func, mat_input)
    assert j.shape == (2, 2, 4, 4)

    在此前的例子中,我们已经介绍了,autograd.backward()为节约空间,仅会保存叶节点的梯度。若我们想得知输出关于某一中间结果的梯度,我们可以选择使用autograd.grad()接口,或是使用hook机制:

    def variable_hook(grad): # hook注册在Tensor上,输入为反传至这一tensor的梯度
    print(‘the gradient of C is:’, grad)

    A = torch.tensor(2., requires_grad=True)
    B = torch.tensor(.5, requires_grad=True)
    C = A * B
    hook_handle = C.register_hook(variable_hook) # 在中间变量C上注册hook
    D = C.exp()
    D.backward() # 反传时打印:the gradient of C is:tensor(2.7183)
    hook_handle.remove() # 如不再需要,可remove掉这一hook

    torch.autograd.gradcheck (数值梯度检查)

    在编写好自己的 autograd function 后,可以利用gradcheck中提供的gradcheckgradgradcheck接口,对数值算得的梯度和求导算得的梯度进行比较,以检查backward是否编写正确。以函数  为例,数值法求得  点的梯度为:  。在下面的例子中,我们自己实现了Sigmoid函数,并利用gradcheck来检查backward的编写是否正确。

    @staticmethod
    def forward(ctx, x):
    output = 1 / (1 + torch.exp(x))
    ctx.save_for_backward(output)
    return output

    @staticmethod
    def backward(ctx, grad_output):
    output, = ctx.saved_tensors
    grad_x = output * (1 output) * grad_output
    return grad_x

    test_input = torch.randn(4, requires_grad=True) # tensor([-0.4646, -0.4403, 1.2525, -0.5953], requires_grad=True)
    torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-3) # pass
    torch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-3) # pass
    torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-4) # fail
    torch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-4) # fail

    我们发现:eps 为 1e-3 时,我们编写的 Sigmoid 和 torch 自带的 builtin Sigmoid 都可以通过梯度检查,但 eps 下降至 1e-4 时,两者反而都无法通过。而一般直觉下,计算数值梯度时, eps 越小,求得的值应该更接近于真实的梯度。这里的反常现象,是由于机器精度带来的误差所致:test_input的类型为torch.float32,因此在 eps 过小的情况下,产生了较大的精度误差(计算数值梯度时,eps 作为被除数),因而与真实精度间产生了较大的 gap。将test_input换为float64的 tensor 后,不再出现这一现象。这点同时提醒我们,在编写backward时,要考虑的数值计算的一些性质,尽可能保留更精确的结果。

    torch.autograd.gradcheck(Sigmoid.apply, (test_input,), eps=1e-6) # pass
    torch.autograd.gradcheck(torch.sigmoid, (test_input,), eps=1e-6) # pass

    torch.autograd.anomaly_mode (在自动求导时检测错误产生路径)

    可用于在自动求导时检测错误产生路径,借助with autograd.detect_anomaly(): 或是 torch.autograd.set_detect_anomaly(True)来启用:

    torch.autograd.grad_mode (设置是否需要梯度)

    我们在 inference 的过程中,不希望 autograd 对 tensor 求导,因为求导需要缓存许多中间结构,增加额外的内存/显存开销。在 inference 时,关闭自动求导可实现一定程度的速度提升,并节省大量内存及显存(被节省的不仅限于原先用于梯度存储的部分)。我们可以利用grad_mode中的troch.no_grad()来关闭自动求导:

    net = resnet50().cuda(0)
    num = 128
    inp = torch.ones([num, 3, 224, 224]).cuda(0)
    net(inp) # 若不开torch.no_grad(),batch_size为128时就会OOM (在1080 Ti上)

    net = resnet50().cuda(1)
    num = 512
    inp = torch.ones([num, 3, 224, 224]).cuda(1)
    with torch.no_grad(): # 打开torch.no_grad()后,batch_size为512时依然能跑inference (节约超过4倍显存)
    net(inp)

    model.eval()torch.no_grad()

    这两项实际无关,在 inference 的过程中需要都打开:model.eval()令 model 中的BatchNormDropout等 module 采用 eval mode,保证 inference 结果的正确性,但不起到节省显存的作用;torch.no_grad()声明不计算梯度,节省大量内存和显存。

    torch.autograd.profiler (提供function级别的统计信息)

    x = torch.randn((1, 3, 224, 224), requires_grad=True)
    model = resnet18()
    with torch.autograd.profiler.profile() as prof:
    for _ in range(100):
    y = model(x)
    y = torch.sum(y)
    y.backward()
    # NOTE: some columns were removed for brevity
    print(prof.key_averages().table(sort_by=“self_cpu_time_total”))

    输出为包含 CPU 时间及占比,调用次数等信息(由于一个 kernel 可能还会调用其他 kernel,因此 Self CPU 指他本身所耗时间(不含其他 kernel 被调用所耗时间)):

    Reference

    [1] Automatic differentiation package – torch.autograd — PyTorch 1.7.0 documentation

    [2] Autograd

    本站上原创文章未经作者许可,不得用于商业用途,仅做学习交流使用,本站免责声明。转载请注明出处,否则保留追究法律责任的权利。《署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0)》许可协议授权
    数据科学与编程 » PyTorch 源码解读之 torch.autograd

    发表评论

    • 52会员总数(位)
    • 307资源总数(个)
    • 40本周发布(个)
    • 1 今日发布(个)
    • 329稳定运行(天)

    提供最优质的博文资源集合

    立即阅览 了解详情