知行编程网知行编程网  2022-07-24 11:00 知行编程网 隐藏边栏 |   抢沙发  120 
文章评分 0 次,平均分 0.0

PyTorch 源码解读之 torch.autograd

作者丨OpenMMLab@知乎

来源丨https://zhuanlan.zhihu.com/p/321449610

编辑丨极市平台

导读

 

本文主要为介绍pytorch中的autograd模块功能为主,包括了autograd的七个模块主要涉及 torch/autograd 下代码,不涉及底层的 C++ 实现。

前言

本篇笔记以介绍 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.Conv2d, nn.ReLU等)作为基本单元。而这些 Module 通常是包裹 autograd function,以其作为真正实现的部分。例如nn.ReLU 实际使用torch.nn.functional.reluF.relu):

这里的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 为例进行讲解:

inplace)

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


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

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

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

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

PyTorch 源码解读之 torch.autograd

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


PyTorch 源码解读之 torch.autograd

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

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


此前的例子使用的是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。

)

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

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

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

我们发现: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.anomaly_mode (在自动求导时检测错误产生路径)

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

 backward

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

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

    net(inp)


model.eval()torch.no_grad()

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

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

))

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

Self CPU time total: 9.217s


Reference

[1] Automatic differentiation package - torch.autograd — PyTorch 1.7.0 documentation(https://pytorch.org/docs/stable/autograd.html)
[2] Autograd(https://blog.csdn.net/qq_43328040/article/details/108421469)


<pre style="font-size: 16px;background-color: rgb(255, 255, 255);text-align: left;color: rgb(86, 86, 86);letter-spacing: 1px;"><section data-brushtype="text" style="padding-right: 0em;padding-left: 0em;white-space: normal;letter-spacing: 0.544px;color: rgb(62, 62, 62);font-family: "Helvetica Neue", Helvetica, "Hiragino Sans GB", "Microsoft YaHei", Arial, sans-serif;widows: 1;word-spacing: 2px;caret-color: rgb(255, 0, 0);text-align: center;"><strong style="color: rgb(0, 0, 0);font-family: -apple-system-font, system-ui, "Helvetica Neue", "PingFang SC", "Hiragino Sans GB", "Microsoft YaHei UI", "Microsoft YaHei", Arial, sans-serif;letter-spacing: 0.544px;"><span style="letter-spacing: 0.5px;font-size: 14px;"><strong style="font-size: 16px;letter-spacing: 0.544px;"><span style="letter-spacing: 0.5px;">—</span></strong>完<strong style="font-size: 16px;letter-spacing: 0.544px;"><span style="letter-spacing: 0.5px;font-size: 14px;"><strong style="font-size: 16px;letter-spacing: 0.544px;"><span style="letter-spacing: 0.5px;">—</span></strong></span></strong></span></strong></section><pre><pre><section style="letter-spacing: 0.544px;white-space: normal;font-family: -apple-system-font, system-ui, "Helvetica Neue", "PingFang SC", "Hiragino Sans GB", "Microsoft YaHei UI", "Microsoft YaHei", Arial, sans-serif;"><section powered-by="xiumi.us"><section style="margin-top: 15px;margin-bottom: 25px;opacity: 0.8;"><section><section style="letter-spacing: 0.544px;"><section powered-by="xiumi.us"><section style="margin-top: 15px;margin-bottom: 25px;opacity: 0.8;"><section><section style="margin-bottom: 15px;padding-right: 0em;padding-left: 0em;color: rgb(127, 127, 127);font-size: 12px;font-family: sans-serif;line-height: 25.5938px;letter-spacing: 3px;text-align: center;"><span style="color: rgb(0, 0, 0);"><strong><span style="font-size: 16px;font-family: 微软雅黑;caret-color: red;">为您推荐</span></strong></span></section><p style="margin: 5px 16px;padding-right: 0em;padding-left: 0em;font-family: sans-serif;letter-spacing: 0px;opacity: 0.8;line-height: normal;text-align: center;">一个算法工程师的日常是怎样的?</p><p style="margin: 5px 16px;padding-right: 0em;padding-left: 0em;font-family: sans-serif;letter-spacing: 0px;opacity: 0.8;line-height: normal;text-align: center;">25个常用Matplotlib图的Python代码<br  /></p><section style="margin-top: 5px;margin-bottom: 5px;padding-right: 0em;padding-left: 0em;min-height: 1em;font-family: sans-serif;letter-spacing: 0px;opacity: 0.8;line-height: normal;text-align: center;"><span style="font-size: 14px;">13个算法工程师必须掌握的PyTorch Tricks</span></section><section style="margin-top: 5px;margin-bottom: 5px;padding-right: 0em;padding-left: 0em;min-height: 1em;font-family: sans-serif;letter-spacing: 0px;opacity: 0.8;line-height: normal;text-align: center;"><span style="font-size: 14px;">吴恩达上新:生成对抗网络(GAN)专项课程</span></section><section style="margin-top: 5px;margin-bottom: 5px;padding-right: 0em;padding-left: 0em;min-height: 1em;font-family: sans-serif;letter-spacing: 0px;opacity: 0.8;line-height: normal;text-align: center;">从SGD到NadaMax,十种优化算法原理及实现</section></section></section></section></section></section></section></section></section>

PyTorch 源码解读之 torch.autograd


本篇文章来源于: 深度学习这件小事

本文为原创文章,版权归所有,欢迎分享本文,转载请保留出处!

知行编程网
知行编程网 关注:1    粉丝:1
这个人很懒,什么都没写

发表评论

表情 格式 链接 私密 签到
扫一扫二维码分享