PyTorch源码分析（2）——动态图原理

[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之间的操作（算子）`两个方面。本节主要关注动态图的原理，因此主要从__算子__层面解析多设备的支持。 我们来看一个例子： {:class="width_60"} 在这个例子中，对于`a.add(b)`这个算子，无论是 _cpu设备_ 的tensor还是 _gpu设备_ 的tensor，都可以得以支持。
__那么同一个add算子是如何支持多设备的呢？__
这里就不得不说PyTorch的__Dispatch机制__。 {:class="width_50"} {: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`，并绑定一个实现函数即可： {:class="width_50"} > PS. 上述的add实现偷了个懒，将tensor先转回cpu，然后算完add再转回device。 然后将文件导出成动态库`libmyop.so`，就可以在pytorch中直接使用了： {:class="width_50"} 可以看到输出了my_add_func，因此调用的是我们自己定义的函数实现。 除了m.def以及m.impl之外，还有__m.fallback__作为回退： {: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`： {:class="width_60"} 该文件位于`aten/src/ATen/native`目录下。 其中每一个算子都以`- func:`开头，而dispatch部分对应的是其具体实现的函数名称。 例如上述的`dot算子`实现，CPU的实现函数名称是`at::native::dot()`，我们可以在CPU算子中找到这个函数实现： {:class="width_60"} 可以看到这里具体实现`at::native::dot()`可以调用`intel MKL库`实现，也有naive的实现（具体看`dot_impl`里面是循环乘法）。
除此之外，在源码部分__不需要__针对性得去写`def`、`impl`等实现。这部分具体是__由codegen模块自动生成__的。 我们在手动编译PyTorch时，可以在`build文件夹`找到自动生成的代码： {:class="width_60"}{:class="width_60"} 而绑定的函数其实就是对前面具体实现`at::native::dot()`的一层封装。
除此之外，在`python tensor`这边，也会__对Tensor注册__`dot()`这个算子函数。这里也是通过算子配置文件利用`codegen`在编译时自动生成的，文件是`torch/csrc/autograd/generated/python_variable_methods.cpp`： {:class="width_60"} 这里`Python<-->c++`需要经过`pybind`的封装，最后绑定到`torch.Tensor`： {:class="width_60"} 整个流程总结如下图：  即PyTorch中只需要实现`native_functions.yaml`这个算子配置文件，以及算子具体实现的函数。其他的都可以利用`codegen模块`自动完成整个流程的注册实现。最终将实现`Tensor的算子成员函数`，供用户使用。 ## 3. 反向传播 ### 1. 反向算子及其注册 类似于前向算子的配置文件`native_functions.yaml`，__反向算子也有一个配置文件`derivatives.yaml`__，其位于`tools/autograd/derivatives.yaml`，我们打开看这个文件大致是这样的： {:class="width_60"} 可以看到，每一个算子以`- name:`开头。 然后还包含一个result字段，这个字段其实就是这个算子的求导公式。 我们知道对dot(a, b)求导的结果是`a'b + b'a`，上图的result即对应这个公式。  前面我们也介绍了前向算子会利用`codegen`自动生成注册部分的代码。同理，__反向算子也可以根据算子微分注册表自动生成dispatch注册，然后被绑定到Python的函数中__（如DotBackward0()）。 而我们在进行PyTorch训练时，对Tensor设置有`requires_grad`属性时，其后续节点都会包含`grad_fn`，即反向函数。`DotBackard0()`就是一种反向函数。 ### 2. PyTorch反向计算原理 > 为弄清前面所说的`grad_fn`以及`反向传播`的计算原理，我们举一个例子来说明。 假设我们Tensor a是一个参数，需要进行求导，loss的公式如下： {:class="width_30"} 将loss公式展开，并画出流程图： {:class="width_45"}{:class="width_45"} 如果我们进行中间过程调试，其实可以看到，在每一条语句完成执行时，该中间变量都会生成一个`grad_fn`成员变量。 {:class="width_45"}{: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`形成了__反向的图结构__。 具体如下： {:class="width_45"} 我们可以手动验证一下利用`grad_fn`来进行求导： {: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`)自动生成的。__ 而算子直接之所以可以通过图的顺序依次通过反向算子计算，其主要是由链式法则决定的： {: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算子总结为下面的图：  在我们调用合成算子的前向时，实际是被拆成多个基础算子被执行的，因此生成的反向图也是基础算子的反向。 实际上PyTorch中__绝大多数算子都属于合成算子__，这也就导致基础算子（需要手动实现反向公式）的算子实际只有__600多个__。 ## 4. 动态图执行过程 ### 1. 前向 对于前向传播来说，我们前面也介绍过。 {:class="width_50"} 实际上在每执行一条python代码时，前向传播的算子都会被实时调用执行。 具体的流程如下图总结：  在用户调用某算子（例如`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++部分。如下图所示： {:class="width_60"} 在c++部分，具体backward()执行实现在`torch/csrc/autograd/engine.cpp`中。 这里首先会找到图的根节点（也就是loss.grad_fn及其参数等），然后执行`execute_with_graph_task`。 {:class="width_60"} 在`execute_with_graph_task()`中，首先会依据设备数建立线程池。 然后root节点会被push到一个__队列__中(queue->push())，然后开启线程（add_trad_pool_task()），后续__线程__执行。 {:class="width_60"} 线程中会调用`thread_main()`，用于开始具体执行。 在`thread_main()`中，会循环判断graph_task->future_result是否完成，即当前线程所有task完成才会阻塞退出返回结果。 我们可以看到在thread_main中，具体就是从队列中取出task，然后进行执行evaluate_function，循环这个过程。 {:class="width_60"} {:class="width_60"} 最终调用的内容是fn(inputs)，也就是`grad_fn(loss)`。 {:class="width_60"} 在执行outputs=call_function()结束后，还会通过遍历output，来获取后续task节点（next_functions）。 对于后续task，会判断该节点是否ready（也就是前置节点是否完成），如果ready了就会被放到ready queue中，供后续调度线程执行。 {:class="width_60"} 上述具体流程总结如下图：  最后总结来说，__在反向传播过程中，会按照图的顺序，多线程调度执行`tensor.grad_fn()`，并发射到device上来实现执行逻辑的__。 ## 5. 总结 最后，用一张图总结一下本文的主要内容。本文主要介绍了PyTorch动态图的原理，主要分为三个部分。  对于__多设备支持__，PyTorch通过`dispatch机制`实现了不同设备下的`算子映射`。PyTorch提供了一套算子注册方法，并在源码中利用`codegen`模块依据算子注册表native_functions.yaml`自动生成算子的注册`。 对于__反向传播__，PyTorch的自动微分模块也是依赖`codegen模块`和`微分注册表(derivatives.yaml)`自动生成反向算子的实现、注册。微分注册表中主要是各类算子的反向公式，并且只包含最基础的算子的反向公式。 对于__动态图原理__，`前向`是在算子调用的时候`直接执行`的算子实现；而`反向`则是通过一套复杂的`Execution Engine`，利用线程、队列等，依次调用执行反向图中的grad_fn（也即反向算子）实现的执行过程。