需要了解的语义和工具
| 语言 | 工具或库 |
|---|---|
| C | CUDA |
| C++ | CMake |
| python | BLAS |
源码中常用的变量缩写

核心代码目录: /XXX/XXX/pytorch/torch
核心模块:
-
TH: 核心就THStorage与THTensor,辅助有THGeneral、THHalf、THAllocator、THSize
-
THS: 稀疏矩阵THSTensor(THLongTensor *indices; THTensor *values;)
-
THNN: C代码的神经网络,每种操作都前向与后向计算,以Abs为例
-
WITH_CUDA: TH、THS、THNN的GPU版本
-
WITH_NCCL: 单机上多GPU通信,支持MPI编程。单PCIe上性能最好,虽然多PCIe也可运行。
-
build THPP: Torch C++
-
build libshm: 共享管理
-
nanopb: 轻量级encode/decode使用Gooogle的Protocol Buffers协议来定义的结构体,以便在RAM或者代码空间中传输消息
-
pybind11: 轻量级的C++与Python间进行无缝调用
-
build ATen: C++的autograd在往这里写,整体功能不太清楚
-
generate_nn_wrappers():把C代码写好的神经网络(NN、CUNN),使用C++进行封装,python能调用
-
C: 对C的代码进行封装,自身是C++的代码,核心为Tensor、autograd、cuda,其它有distributed、jit等;
- torch/csrc/Module.cpp封装成torch模块的方法
- torch/csrc/generic/Tensor.cpp封装成FloatTensorBase、IntTensorBase等类
-
DL: python可获取RTLD_GLOBAL、RTLD_NOW变量,来设置加载库的方式
-
THNN: torch/csrc/nn/THNN.cpp,NN的封装(generate_nn_wrappers函数生成)
-
THCUNN: torch/csrc/nn/THCUNN.cpp,CUNN的封装(generate_nn_wrappers函数生成)
-
from torch._C import *: 加载C实现的所有内容(其实是用C++再封装了一层)
需要了解的接口:
-
BLAS/LAPACK的接口
-
CUBLAS/Thrust:这个是CUDA的BLAS部分会用到。这个文档直接读NVIDIA官方的就好了
-
CPU张量运算库TH
这个部分是PyTorch里最底层的部分,但是也是我觉得写的最好的部分。感觉读懂这一部分其他底层的代码就都很容易读了。它用宏实现了C里的泛型支持,然后用宏+命名实现了简单的OOP。甚至在管理文件的部分还有virtual table。
这个最早是从Lua版本的Torch迁移过来的,那会儿似乎还没有完整的SIMD。现在提供了完整的SIMD支持。数据结构部分主要实现了THStorage和THTensor两个。这两个是干嘛用的在PyTorch的Python文档里都有写。其实比起我们物理这边用的Tensor Network的实现要简单,因为不用连腿....
THTensor不负责存储,是一种查看存储对象(THStorage)的方法。感觉实现思路上和MXNet的backend,mshadow里的Tensor有一些类似的地方。不过因为C不支持重载,所以这里完全没有懒惰求值,当然用的时候也许会有些麻烦。
这里比较有特点的地方是用宏实现的泛型。这个方法在底层的库里有广泛的使用。等到lib目录上层的代码能够使用THPP这个支持C++模板的封装以后就不会见到了。(THPP和facebook/THPP那个不一样,后者是一个C++的TH实现)。
TH主要通过宏替换,比如real(这个名字在我改成复数的时候可别扭了...然后还得小心别和cpp里的std::real一个不小心重名,然后被替换成std::double之类的东西...)真正用来实现的需要泛型的源代码写在generic目录下,然后在外面通过几个头文件里定义的宏进行不同变量名称的替换。
然后还有一个比较难读的宏应该就是TH_TENSOR_APPLY这个开头的几个宏函数了。这几个比较像是在现在更上层语言里用到的map,或者broadcast。也就是把一段操作应用到每个Tensor的元素上去。为了支持不同个数的Tensor之间的运算,有好几个不同个数Tensor的。比如TH_TENSOR_APPLY_3就是说三个Tensor里要怎么做element-wise运算。然后需要注意的是,在宏函数里输入的操作部分也就是CODE部分,一定不要出现逗号,和内嵌的宏。
然后TensorMath和Conv等部分的函数都和Python函数功能类似,不懂直接看Python文档就行了。
CPU神经网络库THNN
这个部分就很简单了,就是一些Python层的Function对象在底层的定义和实现。不懂的就直接查Python文档,都有对应的。
自动微分部分
这个是在C++里写的,然后接了Python。C++里就是实现了三个list用来存Variable啊什么的...大致思路应该是把使用过forward的函数记下来,然后后面调用backward来算梯度。所以叫Tape-based
ps:
结构上大约有两部分,一部分是Python,另外一部分是C/C++写的后端。Python的部分比较容易读,如果后面有时间再介绍。C/C++后端的多维数组库在v0.4开始用aten(a tensor library)库作为上层封装再通过torch/csrc 中的部分胶水代码接入Python(但慢慢在改为pybind11)。自动微分使用C++完成,后端采用的是一个类似于录像带的机制来记录相关操作。
后端核心代码主要在 lib 里