A random walk in PyTorch (2) -- Build the World


按照我本来的计划,第二篇文章应该直接开始看 ATen 了,而且直奔 BLAS 而去。

不过从 NEKit 的经验中我发现,可能是因为大家对于一个新的库都要花不少的时间才能把握住其中的逻辑,所以都不愿意花太多的精力。个人意见,通读代码没什么特别大的必要,除非你要做二次开发。但是至少要掌握库本身的逻辑所在。一个简单的标准就是,如果我问你,XX 是在哪里做的,你至少应该能很快的找到具体的代码位置(虽然你不知道它的具体实现)。靠什么找到呢?靠的是你理解每个模块的功能,所以你自然知道要去找哪一个模块,你也自然明白你的每一行调用究竟做了什么,而不是复制黏贴 demo,再通过随机采样实现自己想要的功能。

我还不太确定,但我估计我之后的各章都只是简单的展开主要逻辑(太懒)。由于我非常可能会省略很多的中间过程,所以可能有的时候会有人会问,你怎么会找到、看到这些东西呢?比如说,为什么会先从 ATen 开始呢?很多事情,自己做很简单,看别人的就很难了。这也是为什么有些时候读书或者读代码很难的原因,你只能看到漫长累加的结果,而不是过程。可能每一个小迭代的逻辑过程都很简单,就好比数学考试的大题若是有五小问就比只有最后一问要容易很多一样,但是你知道自己的逻辑,却很难知道他人的逻辑。所以我想可以先分享一下我读代码的一些思路和逻辑,不成体系,正好借着了解 PyTorch 结构的这个例子来展开一下。

和第一篇一样,这将是一篇不涉及任何 PyTorch 具体技术细节的文章。可以直接跳过。

之后所有内容都基于 commit dd5c195646b941d3e20a72847ac48c41e272b8b2。嗯,因为我拖了很久,其实 master 已经有了一些比较大的变化,比如 Caffe 的代码被合并进了 PyTorch(虽然我还没看到有什么实质的变化),但是整个框架的核心思想没有发生什么改变。另一方面,由于很多原因(比如太懒、太懒和太懒),我不想在花时间比较新的代码和之前的细微变化了,所以一切都已这个 commit 为准。我可能会在全部写完之后再写写一些大的变化(如果有的话)。

Oops, a big project

现在你面对着 PyTorch 的源代码:

-----------------------------------------------------------------------------------
Language                         files          blank        comment           code
-----------------------------------------------------------------------------------
C++                                484          39495          16203         241502
C/C++ Header                       610          12415          15851          83722
Python                             483          19631          21796          64968
C                                  208           6986           3058          40491
CUDA                               315           7128           3859          39864
Markdown                            76           8563              0          30437
YAML                                15            507            199          26066
CMake                               93           1404           2193           7029
MSBuild script                      11              0              0           1711
Protocol Buffers                    58            482            707           1612
Fortran 90                          12            405            161           1528
JSON                                10            131              0           1261
make                                18            273            227            989
Bourne Shell                        26            166            188            793
m4                                   4             68             49            534
CSS                                  3             69             16            284
DOS Batch                            4             42              1            179
Windows Resource File                1              3              0             34
vim script                           2             14              7             34
Dockerfile                           1              7              2             32
Lua                                  1              5              0             28
INI                                  2              0              0             19
HTML                                 1              3              0             12
-----------------------------------------------------------------------------------
SUM:                              2438          97797          64517         543129

从何开始呢?

你对 PyTorch 的依赖组件一无所知,也不了解 PyTorch 代码的组织结构。就算能猜个大概,但怎么把 PyTorch 代码的组织逻辑找出来呢?最起码的,先读哪个文件夹下面的哪个文件?

第一步,也是最重要的,先读完 PyTorch 的文档。读完文档,就应当对一个库的基本功能有了清晰的概念,大体上也可以分析出库的代码有哪些模块所构成。

接下来,你可以看看这个库是怎么编译的(对于某些提供 Makefile 或是 m4 的库就要看情况了),来理解代码的构成(至少你也要能分清楚每个文件夹下面的代码是干什么的)。幸运的是(相对 TensorFlow),PyTorch 的根目录下就有 setup.py,显然 PyTorch 的编译过程不会太难把握。

setup.py 文件的关键部分没有什么技术难度,即便你对 distutilssetuptools 不太了解应该也没啥问题。

我们主要看看 PyTorch 的编译步骤。

Building PyTorch

Oops,setup.py 有 743 行,要看一会儿了……吗?

其实不用,几分钟就看完了。记住你的目的不是学习 setuptools,不要浪费时间,等你要用到的时候自然会弄明白的。

打开文件,先通览一遍,前面好多东西林林总总似乎都是编译的细节,但是在最后:

if __name__ == '__main__':
    setup(name="torch", version=version,
          description="Tensors and Dynamic neural networks in Python with strong GPU acceleration",
          ext_modules=extensions,
          cmdclass=cmdclass,
          packages=packages,
          package_data={'torch': [
              'lib/*.so*', 'lib/*.dylib*', 'lib/*.dll', 'lib/*.lib',
              'lib/torch_shm_manager',
              'lib/*.h',
              'lib/include/TH/*.h', 'lib/include/TH/generic/*.h',
              'lib/include/THC/*.h', 'lib/include/THC/generic/*.h',
              'lib/include/ATen/*.h',
          ]},
          install_requires=['pyyaml', 'numpy'],
          )

好的,我们找到入口了,看一下,ext_modules=extensions 显然是要编译的模块,cmdclass=cmdclass 似乎是可以提供一些命令,packages=packages 显然是一些需要的包。

往上一扫,就在同一屏幕里就有:

cmdclass = {
    'build': build,
    'build_py': build_py,
    'build_ext': build_ext,
    'build_deps': build_deps,
    'build_module': build_module,
    'develop': develop,
    'install': install,
    'clean': clean,
}

大家 python 肯定写的比我多,python setup.py install 大家肯定是见过的,那么我们可以看看这里 install 是在干什么。

class install(setuptools.command.install.install):

    def run(self):
        if not self.skip_build:
            self.run_command('build_deps')
        setuptools.command.install.install.run(self)

嗯……默认实现外加运行 build_deps。但是默认实现是什么呢?看起来有一点麻烦。

但是,默认实现还能干什么呢,还不就是 build 需要 build 的东西么?需要 build 的东西是什么呢,必然就在 build* 命令里了。

我们看一下:一共五个 buildbuild_py 是默认实现估计干不了啥,build_ext 显然是 build extensions,build_module 调用了 build_pybuild_ext,还有一个 build_deps

显然,我们只需要看看 build_deps 如何处理了哪些依赖和 build_ext 如何如何编译 extension 就可以了。

从另一个角度,setuptools 是一个很完善的库,所以从命令行应该可以获取到一些信息,我们试试

python setup.py -h

就可以知道 --help-commands 命令,再通过

python setup.py --help-commands

就可以得到每个命令的说明。

显然 build 命令是我们唯一需要的。

让我们先试一试,

pip install pyyaml numpy
python setup.py build

Works like a charm.

Why?

为什么不仔细读每一行代码,尤其是,为什么不先读一下 setuptools 的文档呢?

最主要的原因,当然是因为 setuptools 也好,具体编译的 flag 也好,都无关我们的目的。在这里花费时间没有什么意义。

另外,我想要说明一个我阅读代码的大方向和思路。PyTorch 刚好是用 python 写的 setup.py,比较简单,你也可以 argue 说学了也比较实用。若是改天遇上一个手写 Makefile 的怎么办?难不成先学一遍 Makefile 不成?但是你又没有写 Makefile 的需求,也不会对你了解项目有太大帮助,而且这年头你也不太可能在任何新项目中写 Makefile 了,而且学 Makefile 是不是应该先精通 shell?若是你读 TensorFlow,还要先学会 Bazel 么,学会了,你自己的项目用得起来么?

接口也好,语法也罢,都是无关痛痒的,今天看了一遍,明天就忘了。种种 tools,追根究底只是一个 implementation,就说 Bazel,有意义的问题只有,为什么它那么快(可能还有,为什么你用不起来),至于如何配置,不过是一个某种意义上很随意的选择,除非决定要用,不然没什么学习的价值。说老实话,我连很多我自己一行行写出来的东西都不记得了。看其他的项目,我觉得最有价值的是学到正确的做法是什么。比如说对于这个 setup.py,如果通过上面的阅读,理解 setup.py 的基本工作模式和原理,并且知道在 python 中如果开发一个库应该使用 setup.pysetuptools (或者 distutils)来打包分发就可以了。最好还可以记住 PyTorch 用到了 setuptools 比较全面的功能,可以在需要的时候拿来作为 demo 参考,我觉得就足够了。

What’s in there

好吧,那么依赖有哪些呢,ATen,nanopb,libshm,gloo,THD,nccl,调用 torch/lib/build_libs.sh 来编译。

torch/lib/build_libs.sh 就很清楚了,利用 CMake 编译这些依赖,现在还没必要仔细研究编译过程。

重要的是这些依赖的功能。

ATen,核心中的核心,PyTorch 的 Tensor 库。

nanopb,看起来是一个精简版的 protobuf,PyTorch 需要它做什么呢?还不清楚。不过可以想见,应当可以用在多机并行和保存(序列化)模型的过程中。暂时不需要深究这些问题,重要的是我们要记住有这个库可以做这些事就可以了。

libshm,共享内存。我的第一反应是用在 DataLoader 或者类似的情况中,在多进程间共享 Tensor。Again,我写这一系列文章的时候对于 PyTorch 代码的理解并不比大家多,所以我们可以一起来慢慢看。

gloo,并行算法库。

THD,从名字来看应当是并行时用到的 Tensor 库,不过信息太少,我们之后再看。

另外可以看到 torch/lib 下面的 pybind11。用来绑定 python 和 C++ 的结构的库。如果读一下 README,它是 header only 的,再在项目中搜索一下 pybind 就会看到它是在编译 PyTorch 的 extension 的时候编译的。

好,那么大体上我们已经把整个结构理清了。

How PyTorch is organized

aten: Tensor 库。顺带一说,虽然理论上是说 Torch 和 PyTorch 要齐头并进,但是事实上 Torch 已经快半年没更新了,到现在 Torch 也没用上新的 Aten ……

cmake: 一些 CMake 相关的脚本。

docs: 文档,在官网上读就可以了。

test: 可以略过。

toolstorch …… 嗯……

我们再回过头来看看 setup.py 里的 packages,可以看到

packages = find_packages(exclude=('tools', 'tools.*',))

最终安装的 python 代码只包括 torch 下的的文件。显然, tools 只是为了开发和编译而存在的。

tools: 开发和编译中的工具。

torch: PyTorch 的源代码。

Let’s get the party started

好,我们可以正式的开始阅读代码了。代码主要分布在 atentorch 两个文件夹中。首先,我们就从 PyTorch 的最核心,最根本的部分,Aten 开始。