深入浅出PyTorch(算子篇)

Tensor

自从张量(Tensor)计算这个概念出现后,神经网络的算法就可以看作是一系列的张量计算。所谓的张量,它原本是个数学概念,表示各种向量或者数值之间的关系。PyTorch的张量(torch.Tensor)表示的是N维矩阵与一维数组的关系。

http://web.mit.edu/~ezyang/Public/pytorch-internals.pdf

torch.Tensor的使用方法和numpy很相似(https://pytorch.org/...tensor-tutorial-py),两者唯一的区别在于torch.Tensor可以使用GPU来计算,这就比用CPU的numpy要快很多。

张量计算的种类有很多,比如加法、乘法、矩阵相乘、矩阵转置等,这些计算被称为算子(Operator),它们是PyTorch的核心组件。

算子的backend一般是C/C++的拓展程序,PyTorch的backend是称为"ATen"的C/C++库,ATen是"A Tensor"的缩写。

Operator

PyTorch所有的Operator都定义在Declarations.cwrap和native_functions.yaml这两个文件中,前者定义了从Torch那继承来的legacy operator(aten/src/TH),后者定义的是native operator,是PyTorch的operator。

相比于用C++开发的native code,legacy code是在PyTorch编译时由gen.py根据Declarations.cwrap的内容动态生成的。因此,如果你想要trace这些code,需要先编译PyTorch。

legacy code的开发要比native code复杂得多。如果可以的话,建议你尽量避开它们。

aten/src/ATen/Declarations.cwrap

MatMul

本文会以矩阵相乘--torch.matmul()为例来分析PyTorch算子的工作流程。

我在深入浅出全连接层(fully connected layer)中有讲在GPU层面是如何进行矩阵相乘的。Nvidia、AMD等公司提供了优化好的线性代数计算库--cuBLAS/rocBLAS/openBLAS,PyTorch只需要调用它们的API即可。

Figure 1: function flow of torch.matmul()

Figure 1是torch.matmul()在ATen中的function flow。可以看到,这个flow可不短,这主要是因为不同类型的tensor(2d or Nd, batched gemm or not,with or without bias,cuda or cpu)的操作也不尽相同。

at::matmul()主要负责将Tensor转换成cuBLAS需要的格式。前面说过,Tensor可以是N维矩阵,如果tensor A是3d矩阵,tensor B是2d矩阵,就需要先将3d转成2d;如果它们都是>=3d的矩阵,就要考虑batched matmul的情况;如果bias=True,后续就应该交给at::addmm()来处理;总之,matmul要考虑的事情比想象中要多。

除此之外,不同的dtype、device和layout需要调用不同的操作函数,这部分工作交由c10::dispatcher来完成。

Dispatcher

dispatcher主要用于动态调用dtype、device以及layout等方法函数。用过numpy的都知道,np.array()的数据类型有:float32, float16,int8,int32,.... 如果你了解C++就会知道,这类程序最适合用模板(template)来实现。

很遗憾,由于ATen有一部分operator是用C语言写的(从Torch继承过来),不支持模板功能,因此,就需要dispatcher这样的动态调度器。

类似地,PyTorch的tensor不仅可以运行在GPU上,还可以跑在CPU、mkldnn和xla等设备,Figure 1中的dispatcher4就根据tensor的device调用了mm的GPU实现。

layout是指tensor中元素的排布。一般来说,矩阵的排布都是紧凑型的,也就是strided layout。而那些有着大量0的稀疏矩阵,相应地就是sparse layout。

Figure 2: strided layout example

Figure 2是strided layout的演示实例,这里创建了一个2行2列的矩阵a,它的数据实际存放在一维数组(a.storage)里,2行2列只是这个数组的视图。

stride充当了从数组到视图的桥梁,比如,要打印第2行第2列的元素时,可以通过公式:\(1 * stride(0) + 1 * stride(1)\)来计算该元素在数组中的索引。

除了dtype、device、layout之外,dispatcher还可以用来调用legacy operator。比如说addmm这个operator,它的GPU实现就是通过dispatcher来跳转到legacy::cuda::_th_addmm。

aten/src/ATen/native/native_functions.yaml

END

到此,就完成了对PyTorch算子的学习。如果你要学习其他算子,可以先从aten/src/ATen/native目录的相关函数入手,从native_functions.yaml中找到dispatch目标函数,详情可以参考Figure 1。


更多精彩文章,欢迎扫码关注下方的公众号, 并访问我的简书博客:https://www.jianshu.com/u/c0fe8671254e

欢迎转发至朋友圈,工作号转载请后台留言申请授权~

AI实战:一个有料有深度的公众号

posted @ 2020-06-22 22:23  A君来了  阅读(7459)  评论(0编辑  收藏  举报