PyTorch源码分析(2)——动态图原理

    Hurray 1825次浏览 4条评论 5609字

#深度学习 #框架 #PyTorch #源码

[TOC] ## 1. 回顾与本文概述 前一篇([https://hurray0.com/menu/151/](https://hurray0.com/menu/151/))中,我们介绍了PyTorch的代码架构、前后端以及动态图、静态图的概述。 本文中,我们将详细介绍PyTorch的__动态图及其原理__。
我们前面说过,__训练框架__最重要的特点是: 1. 支持类似numpy的`张量计算`,可以使用`GPU`加速; 2. 支持带`自动微分系统`的深度神经网络; 此外__PyTorch的特征__还包括原生支持`动态图`执行。
那么依据这三点,我们__针对PyTorch提出三个问题__: 1. PyTorch如何支持CPU、GPU等`诸多设备`的? 2. PyTorch如何实现`自动微分`的? 3. `动态图的原理`是什么? 接下来我们将带着这三个问题进行讲解。 ## 2. 算子支持、Dispatch机制 ### 1. dispatch原理 首先__多设备的支持__包括`tensor的存储`以及`tensor之间的操作(算子)`两个方面。本节主要关注动态图的原理,因此主要从__算子__层面解析多设备的支持。 我们来看一个例子: ![](https://file.hurray0.com/uploads/menu/152/31ba93996fdf88106876e884f6cfdf40.png){:class="width_60"} 在这个例子中,对于`a.add(b)`这个算子,无论是 _cpu设备_ 的tensor还是 _gpu设备_ 的tensor,都可以得以支持。
__那么同一个add算子是如何支持多设备的呢?__
这里就不得不说PyTorch的__Dispatch机制__。 ![](https://file.hurray0.com/uploads/menu/152/02ae6527955d4c2e76f62e6bf28f3661.png){:class="width_50"} ![](https://file.hurray0.com/uploads/menu/152/a9b064f3abb6c56871462aaa83312b90.png){:class="width_50"} 我们可以将`Dispatch机制`看做一个__二维的表结构__。其一个维度是`各类设备`(CPU、CUDA、XLA、ROCM等等),一个维度是`各类算子`(add、mul、sub等等)。 PyTorch提供了一套`定义(def)`、`实现(impl)`机制,可以实现__某算子__在__某设备(dispatch key)__的绑定。 例如上图`m.impl()`中就是对dispatch key为`CPU`时`neg算子`的实现绑定,其绑定了`neg_cpu()`这个函数。
> 举个例子,__如何实现一个cuda下的自定义add算子呢?__ 由于m.def是不限定设备(dispatch key)的共用定义,所以不需要重新定义一遍。我们只需要实现`m.impl`,并绑定一个实现函数即可: ![](https://file.hurray0.com/uploads/menu/152/07c3239f601aad5504c9400cc4aa6269.png){:class="width_50"} > PS. 上述的add实现偷了个懒,将tensor先转回cpu,然后算完add再转回device。 然后将文件导出成动态库`libmyop.so`,就可以在pytorch中直接使用了: ![](https://file.hurray0.com/uploads/menu/152/bb44fb44649b0a2a8585c72804746aa3.png){:class="width_50"} 可以看到输出了my_add_func,因此调用的是我们自己定义的函数实现。 除了m.def以及m.impl之外,还有__m.fallback__作为回退: ![](https://file.hurray0.com/uploads/menu/152/a893154d007eda004c0797ff91e0731d.png){:class="width_50"} 这里的`fallback`是指,在没有m.impl实现的情况下,`默认回退的实现`(例如fallback回cpu实现,即D2H到CPU内存,CPU计算完再H2D回设备内存)。 这样我们将不需要对cuda实现100%的算子实现,而是优先实现高优先级的算子,减少新设备情况下的开发量,而未被实现的算子则默认被fallback实现。 ### 2. PyTorch算子流程 > 清楚了Dispatch机制之后,我们将来看看__PyTorch具体是怎么完成算子这套具体流程的__。 我们清楚类似`add`这样的算子,会被PyTorch绑定到`Python`中的`Tensor成员函数`实现中,这样就可以直接使用`a.add(b)`这样的语句。 那么PyTorch多达2000余个算子如果都手动实现一遍,是一个非常繁琐的过程。 因此,在PyTorch中采用了__算子配置文件__`native_functions.yaml`,配合`codegen`模块__自动完成__整个流程。 我们先来看看算子配置文件`native_functions.yaml`: ![](https://file.hurray0.com/uploads/menu/152/0fcbbbcbdaf350bdab872b423f94a940.png){:class="width_60"} 该文件位于`aten/src/ATen/native`目录下。 其中每一个算子都以`- func:`开头,而dispatch部分对应的是其具体实现的函数名称。 例如上述的`dot算子`实现,CPU的实现函数名称是`at::native::dot()`,我们可以在CPU算子中找到这个函数实现: ![](https://file.hurray0.com/uploads/menu/152/2f1228d52f7fda833df3a198d75c32e0.png){:class="width_60"} 可以看到这里具体实现`at::native::dot()`可以调用`intel MKL库`实现,也有naive的实现(具体看`dot_impl`里面是循环乘法)。
除此之外,在源码部分__不需要__针对性得去写`def`、`impl`等实现。这部分具体是__由codegen模块自动生成__的。 我们在手动编译PyTorch时,可以在`build文件夹`找到自动生成的代码: ![](https://file.hurray0.com/uploads/menu/152/6eb884cb2d9c70a3a58284a968bd946f.png){:class="width_60"}![](https://file.hurray0.com/uploads/menu/152/82d80be3e70c2e751e620006b51d098e.png){:class="width_60"} 而绑定的函数其实就是对前面具体实现`at::native::dot()`的一层封装。
除此之外,在`python tensor`这边,也会__对Tensor注册__`dot()`这个算子函数。这里也是通过算子配置文件利用`codegen`在编译时自动生成的,文件是`torch/csrc/autograd/generated/python_variable_methods.cpp`: ![](https://file.hurray0.com/uploads/menu/152/1a1290b2b398e619855e6a8cfbf45887.png){:class="width_60"} 这里`Python<-->c++`需要经过`pybind`的封装,最后绑定到`torch.Tensor`: ![](https://file.hurray0.com/uploads/menu/152/616a6d31ddc17af3a27c5d70e7df4e77.png){:class="width_60"} 整个流程总结如下图: ![](https://file.hurray0.com/uploads/menu/152/4f6d6dcd564690d324e9a5d569295a98.png) 即PyTorch中只需要实现`native_functions.yaml`这个算子配置文件,以及算子具体实现的函数。其他的都可以利用`codegen模块`自动完成整个流程的注册实现。最终将实现`Tensor的算子成员函数`,供用户使用。 ## 3. 反向传播 ### 1. 反向算子及其注册 类似于前向算子的配置文件`native_functions.yaml`,__反向算子也有一个配置文件`derivatives.yaml`__,其位于`tools/autograd/derivatives.yaml`,我们打开看这个文件大致是这样的: ![](https://file.hurray0.com/uploads/menu/152/e8f4dac13bcc8368349188917ed39782.png){:class="width_60"} 可以看到,每一个算子以`- name:`开头。 然后还包含一个result字段,这个字段其实就是这个算子的求导公式。 我们知道对dot(a, b)求导的结果是`a'b + b'a`,上图的result即对应这个公式。 ![](https://file.hurray0.com/uploads/menu/152/38ec71fa781bda759d4a02a50b28485d.png) 前面我们也介绍了前向算子会利用`codegen`自动生成注册部分的代码。同理,__反向算子也可以根据算子微分注册表自动生成dispatch注册,然后被绑定到Python的函数中__(如DotBackward0())。 而我们在进行PyTorch训练时,对Tensor设置有`requires_grad`属性时,其后续节点都会包含`grad_fn`,即反向函数。`DotBackard0()`就是一种反向函数。 ### 2. PyTorch反向计算原理 > 为弄清前面所说的`grad_fn`以及`反向传播`的计算原理,我们举一个例子来说明。 假设我们Tensor a是一个参数,需要进行求导,loss的公式如下: ![](https://file.hurray0.com/uploads/menu/152/e659f97d1fa01ba783a11824290809a8.png){:class="width_30"} 将loss公式展开,并画出流程图: ![](https://file.hurray0.com/uploads/menu/152/cd012fe9a22ede9e0e2ba86119fa3463.png){:class="width_45"}![](https://file.hurray0.com/uploads/menu/152/b60c082f19cd4dba852e515e56a550f5.png){:class="width_45"} 如果我们进行中间过程调试,其实可以看到,在每一条语句完成执行时,该中间变量都会生成一个`grad_fn`成员变量。 ![](https://file.hurray0.com/uploads/menu/152/e666d0bf6e10c5bedd3f86ed776df962.png){:class="width_45"}![](https://file.hurray0.com/uploads/menu/152/d88f4f04e0cb93fe1eec541b6fc1c340.png){:class="width_45"} 我们仔细查看`grad_fn`的内容其实可以发现,`loss=c.sum()`,而其grad_fn是`SumBackward`,`c=b1*b2`,而其`grad_fn`是`MulBackward`…… 因此,`grad_fn`其实就是该result tensor的反向算子。 除此之外,我们还可以发现,`grad_fn`直接是有关联的。例如`c.grad_`和`loss.grad_fn.next_functions[0][0]`是__等价的__。 也就是说,`grad_fn`之间通过`next_functions`形成了__反向的图结构__。 具体如下: ![](https://file.hurray0.com/uploads/menu/152/86fee16f785230faf46d87d1500fc419.png){:class="width_45"} 我们可以手动验证一下利用`grad_fn`来进行求导: ![](https://file.hurray0.com/uploads/menu/152/3eb4fd87da2a54a5df95fe36f99f2dda.png){:class="width_60"} 其中`dloss`是loss对loss进行求导,因此结果是1。 `dc`是loss对c进行求导,其结果是loss.grad_fn(dloss),依次类推。 我们可以发现按照这种方法计算出的`da`,也就是loss对a进行求导,其结果与执行`loss.backward()`后a.grad的值是相同的。 这也验证了我们推断的正确性。
因此,我们可以总结: 1. __前向计算算子执行时,每个tensor会绑定一个`grad_fn`,也就是反向算子;__ 2. __反向算子是由反向公式表(`derivatives.yaml`)自动生成的。__ 而算子直接之所以可以通过图的顺序依次通过反向算子计算,其主要是由链式法则决定的: ![](https://file.hurray0.com/uploads/menu/152/bceb7e3888a10a3540aee46916fa512c.png){:class="width_30"} ### 3. 反向算子有哪些 实际上,我们对比`算子注册表(native_functions.yaml)`和`反向公式表(derivatives.yaml)`中的算子个数可以发现,前者具有2495个算子,而后者只有642个。 造成二者算子数量差异的主要原因是,__并非所有算子都需要实现其反向__。 我们举一个例子,部分算子例如`linear`,其实可以被拆分成`matmul+add`两个算子实现,那么我只需要实现matmul、add这样的基础算子的反向就可以了。我们将类似这样调用基础算子实现的算子总结为`“合成算子”`。 除此之外,`matmul`实际也可以根据shape的不同,调用不同的矩阵乘算子实现,例如`dot(1x1)`、`mv(2x1)`、`mm(2x2)`等。所以matmul也属于合成算子,不需要单独实现反向公式。 另外还有一些特殊情况,例如 部分算子不存在简单数学公式来计算反向 ,那么PyTorch也可以为其反向单独实现一个函数来绑定;另外还有一些算子不存在反向的意义,例如一些拷贝算子等等。
我们将PyTorch算子总结为下面的图: ![](https://file.hurray0.com/uploads/menu/152/b199359fbc658724f6ac30ad113dcbec.png) 在我们调用合成算子的前向时,实际是被拆成多个基础算子被执行的,因此生成的反向图也是基础算子的反向。 实际上PyTorch中__绝大多数算子都属于合成算子__,这也就导致基础算子(需要手动实现反向公式)的算子实际只有__600多个__。 ## 4. 动态图执行过程 ### 1. 前向 对于前向传播来说,我们前面也介绍过。 ![](https://file.hurray0.com/uploads/menu/152/99383df94e54a7d285240fc86546f075.png){:class="width_50"} 实际上在每执行一条python代码时,前向传播的算子都会被实时调用执行。 具体的流程如下图总结: ![](https://file.hurray0.com/uploads/menu/152/b101bc85d22ad16cec3e3daf41bff389.png) 在用户调用某算子(例如`dot`时),其实调用的是Tensor下的`dot()`函数实现。其具体实现在c++中,经过`pybind`和`dispatch(选择设备)`机制后定位带`at::native::dot()`函数。而后对于CPU来说,可以调用intel MKL库的`mkldnn_matmul()`实现。 ### 2. 反向 对于反向来说,其执行过程相对复杂。 前面章节3.反向传播中介绍到,在执行`loss.backward()`时,实际调用执行的是各中间tensor的`grad_fn`。 由于反向计算时会组成一个由`grad_fn`为节点,`next_functions`为边的__反向图__,因此如何高效执行这个图成为一个问题。 如果采用naive的方法深度优先或广度优先遍历执行,则无法利用图的特点进行并行优化。 因此在__PyTorch中存在一套调度执行机制execution engine__。 在用户执行`loss.backward()`时,实际执行的是`Tensor.backward()`。这个backward()向下溯源最终会定位到`autograd模块`中的`backward()`实现,然后进入c++部分。如下图所示: ![](https://file.hurray0.com/uploads/menu/152/0a44fcb018f2becf1f8ec080e1c3ca5b.png){:class="width_60"} 在c++部分,具体backward()执行实现在`torch/csrc/autograd/engine.cpp`中。 这里首先会找到图的根节点(也就是loss.grad_fn及其参数等),然后执行`execute_with_graph_task`。 ![](https://file.hurray0.com/uploads/menu/152/52502fe1105782107c994dfe79484486.png){:class="width_60"} 在`execute_with_graph_task()`中,首先会依据设备数建立线程池。 然后root节点会被push到一个__队列__中(queue->push()),然后开启线程(add_trad_pool_task()),后续__线程__执行。 ![](https://file.hurray0.com/uploads/menu/152/e975a6ef75dc2bc31bbfdd797e89b9dc.png){:class="width_60"} 线程中会调用`thread_main()`,用于开始具体执行。 在`thread_main()`中,会循环判断graph_task->future_result是否完成,即当前线程所有task完成才会阻塞退出返回结果。 我们可以看到在thread_main中,具体就是从队列中取出task,然后进行执行evaluate_function,循环这个过程。 ![](https://file.hurray0.com/uploads/menu/152/cfa19507ca8bc324b2341ae7bd142ec3.png){:class="width_60"} ![](https://file.hurray0.com/uploads/menu/152/93f12ab3fa2079b5fb239136c7ecb308.png){:class="width_60"} 最终调用的内容是fn(inputs),也就是`grad_fn(loss)`。 ![](https://file.hurray0.com/uploads/menu/152/be588aefee7f7094f0065242f3d9eb01.png){:class="width_60"} 在执行outputs=call_function()结束后,还会通过遍历output,来获取后续task节点(next_functions)。 对于后续task,会判断该节点是否ready(也就是前置节点是否完成),如果ready了就会被放到ready queue中,供后续调度线程执行。 ![](https://file.hurray0.com/uploads/menu/152/85ab88671df4a2f866deb85dde7c040c.png){:class="width_60"} 上述具体流程总结如下图: ![](https://file.hurray0.com/uploads/menu/152/cfff3469d83e94be570a69a266b1dd1e.png) 最后总结来说,__在反向传播过程中,会按照图的顺序,多线程调度执行`tensor.grad_fn()`,并发射到device上来实现执行逻辑的__。 ## 5. 总结 最后,用一张图总结一下本文的主要内容。本文主要介绍了PyTorch动态图的原理,主要分为三个部分。 ![](https://file.hurray0.com/uploads/menu/152/4040113e79259c75ff4edc77cc908135.png) 对于__多设备支持__,PyTorch通过`dispatch机制`实现了不同设备下的`算子映射`。PyTorch提供了一套算子注册方法,并在源码中利用`codegen`模块依据算子注册表native_functions.yaml`自动生成算子的注册`。 对于__反向传播__,PyTorch的自动微分模块也是依赖`codegen模块`和`微分注册表(derivatives.yaml)`自动生成反向算子的实现、注册。微分注册表中主要是各类算子的反向公式,并且只包含最基础的算子的反向公式。 对于__动态图原理__,`前向`是在算子调用的时候`直接执行`的算子实现;而`反向`则是通过一套复杂的`Execution Engine`,利用线程、队列等,依次调用执行反向图中的grad_fn(也即反向算子)实现的执行过程。

最后修改: