【pytorch源码分析 |
您所在的位置：网站首页 › pytorch算子包括哪些 › 【pytorch源码分析 |
【pytorch源码分析
|
模块介绍
aten: A Tensor Library的缩写。与Tensor相关的内容都放在这个目录下。如Tensor的定义、存储、Tensor间的操作（即算子/OP）等 torch: 即PyTorch的前端代码。我们用户在import torch时实际引入的是这个目录。 其中包括前端的Python文件，也包括高性能的c++底层实现（csrc/）。为实现Python和c++模块的打通，这里使用了pybind作为胶水。在python中使用torch._C.[name]实际调用的就是libtorch.so中的c++实现，而PyTorch在前端将其进一步封装为python函数供用户调用 c10、caffe2：移植caffe后端，c10指的是caffe tensor library，相当于caffe的aten。 PyTorch1.0完整移植了caffe2的源码，将两个项目进行了合并。引入caffe的原因是Pytorch本身拥有良好的前端，caffe2拥有良好的后端，二者在开发过程中拥有大量共享代码和库。简而言之，caffe2是一个c++代码，实现了各种设备后端逻辑 tools：用于代码自动生成（codegen），例如autograd根据配置文件实现反向求导OP的映射。 scripts：一些脚本，用于不同平台项目构建或其他功能性脚本 小结 PyTorch源码中，最重要的两个目录是aten和torch目录aten(A Tensor Library)目录主要是和Tensor相关实现的目录，包括算子的具体实现torch目录是PyTorch前端及其底层实现，用户import torch即安装的这个目录 torch前端与后端PyTorch 中，前端指的是 PyTorch 的 Python 接口， 后端指的是 PyTorch 的底层 C++ 引擎，它负责执行前端指定的计算 后端引擎也负责与底层平台（如 GPU 和 CPU）进行交互，并将计算转换为底层平台能够执行的指令
编译后的torch前端接口没有csrc后端接口，该部分c++内容（csrc目录）并没有被复制过来，而是以编译好的动态库文件（_C.cpython-*.so） 前后端交互流程 我们以torch.Tensor为例 import torch torch.Tensor
结论是对应实现在 torch/csrc/autograd/python_variable.cpp中，而这个是通过编译后的so包实现_C调用，因为pyi是一个python存根文件，只有定义没有实现，实现都在python_variable.cpp中 看看对应c++的实现逻
一个主流的训练框架需要有两大特征： 实现类似numpy的张量计算，可以使用GPU进行加速； 实现带自动微分系统的深度神经网络 静态图 在TensorFlow1.x中，我们如果需要执行计算，需要建立一个session，并执行session.run()来执行 import tensorflow as tf a = tf.ones(5) b = tf.ones(5) c = a + b sess = tf.Session() print(sess.run(c)) 其整个过程其实是将计算过程构成了一张计算图，然后运行这个图的根节点。这样先构成图，再运行图的方式我们称为静态图或者图模式 动态图 在PyTorch中，我们可以在计算的任意步骤直接输出结果（当然最新的ensorflow已经支持动态模式）PyTorch每一条语句是同步执行的，即每一条语句都是一个（或多个）算子，被调用时实时执行。这种实时执行算子的方式我们称为动态图或算子模式 import tensorflow as tf a = tf.ones(5) b = tf.ones(5) c = a + b print(c) # tf.Tensor([2. 2. 2. 2. 2.], shape=(5,), dtype=float32) import torch a = torch.ones(5) b = torch.ones(5) c = a + b print(c) # tensor([2., 2., 2., 2., 2.]) 小结 动态图的优点显而易见，可以兼容Python式的编程风格，实时打印编程结果，用户友好性做到最佳；静态图则在性能方面有一定优势，即在整个图执行前，可以将整张图进行编译优化，通过融合等策略改变图结构，从而实现较好的性能PyTorch原生支持动态图，但是也在静态图方面做了诸多尝试，例如 torchscript(jit.script、jit.trace)、TorchDynamo、torch.fx、LazyTensor 图原理
动态图&静态图&dispatch原理解读视频 动态图 首先PyTorch的动态图是从Python源码下降，拆分成多个python算子调用，具体调用到tensor的OP。经过pybind转换到c++，并通过dispatch机制选择不同设备下的算子实现，最终实现调用的底层设备实现（如Nvidia cudnn、intel mkl等） 静态图 jit.script jit.script是在python源码角度分析function的源码，将python code转换为图（torchscript IR格式）。由于是直接从python源码转的，因此有许多python语法无法完备支持，存在转换失败的可能性 jit.trace jit.trace则是在c++层面获取算子，在算子调用时记录成图（torchscript IR格式）。由于需要下降到c++，因此需要输入一遍数据真正“执行”一遍。而获取的图具有确定性，即如果图存在分支，则只能被trace记录其中一个分支的计算路线 torch.fx python层面算子调用时记录的算子，输出是fx IR格式的图 LazyTensor 在算子调用时记录成的图，该调用会截获正常算子的运行，在用户指定同步时再整体运行已积累的图 TorchDynamo 将python源码转变成二进制后，通过分析二进制源码，获取分析的算子和图。最新的PyTorch1.14（2.0）中将其作为torch.compile()的主要路线 算子支持与dispatch机制小结训练框架最重要的特点是： 支持类似numpy的张量计算，可以使用GPU加速；支持带自动微分系统的深度神经网络；原生支持动态图执行针对以上问题，提出3个问题： PyTorch如何支持CPU、GPU等诸多设备的?PyTorch如何实现自动微分的?动态图的原理是什么? 动态图Dispatch机制 import torch a = torch.randn(5, 5) b = torch.randn(5, 5) c = a.add(b) print(c.device) # cpu a = a.to("cuda") b = b.to("cuda") c = a.add(b) print(c.device) # cuda 上述示例中，a.add(b)这个算子，无论是cpu设备的tensor还是gpu设备的tensor，都可以得以支持为什么同一个算子在不同设备上都能运行呢? dispatch机制 文档：https://pytorch.org/tutorials/advanced/dispatcher.html 原理
例如m.impl()中就是对dispatch key为CPU时neg算子的实现绑定，其绑定了neg_cpu()这个函数 大多数情况我们只需要实现m.impl，并绑定一个实现函数即可 除了m.def以及m.impl之外，还有m.fallback作为回退 在没有m.impl实现的情况下，默认回退的实现（例如fallback回cpu实现）。这样我们将不需要对cuda实现100%的算子实现，而是优先实现高优先级的算子，减少新设备情况下的开发量，而未被实现的算子则默认被fallback实现 实现一个定义算子add覆盖原始add算子（todo） 算子配置文件native_functions.yaml PyTorch中采用了算子配置文件aten/src/ATen/native/native_functions.yaml，配合codegen模块自动完成整个流程 也就是多有的自动注册流程会基于当前这个yaml配置文件自动生成算子注册方法与python bind实现 以dot算子为例 配置 算子实现  上诉代码手动编译后，由codegen会自动生成def、impl的实现，也会自动生成pybind的实现build文件夹找到自动生成的代码 pytorch/build/aten/src/ATen/RegisterCPU.cpp torch/csrc/autograd/generated/python_variable_methods.cpp
反向传播
dispatch实际是前向算子 类似的也有反向算子，配置文件derivatives.yaml，其位于tools/autograd/derivatives.yaml  最后利用codegen根据算子微分注册表自动生成dispatch注册，然后被绑定到Python的函数中
动态图执行过程
前向过程
import torch
a = torch.ones(5, 5)
b = torch.ones(5, 5)
c = a + b
print(c)
实际上在每执行一条python代码时，前向传播的算子都会被实时调用执行 在用户调用某算子（例如dot时），其实调用的是Tensor下的dot（）函数实现。其具体实现在c++中，经过pybind和dispatch（选择设备）机制后定位带at::native::dot()函数。而后对于CPU来说，可以调用intel MKL库的mkldnn_matmul()实现 反向过程 import torch a = torch.ones(5, 5, requires_grad=True) b = torch.ones(5, 5, requires_grad=True) c = (a + b).sum() c.backward() print(c) print(c.grad_fn) # tensor(50., grad_fn=) #在执行loss.backward()时，实际调用执行的是各中间tensor的grad_fn，由于反向计算时会组成一个由grad_fn为节点，next_functions为边的反向图，因此如何高效执行这个图成为一个问题 为了解决这个问题，引入了根据设备数建立的线程池调度引擎 总结
|
可以看到在aten/src/Aten目录下，算子实现都在native/目录中。其中有CPU的算子实现，以及CUDA的算子实现（cuda/）等
可以看到，每一个算子以“- name:”开头。 然后还包含一个result字段，这个字段其实就是这个算子的求导公式 前向算子会利用codegen自动生成注册部分的代码。同理，反向算子也可以根据算子微分注册表自动生成dispatch注册，然后被绑定到Python的函数中 有关梯度计算请参考
【本文地址】
今日新闻 |
点击排行 |
|
推荐新闻 |
图片新闻 |
|
专题文章 |