跳转到主要内容
博客

PyTorch 内部结构第二部分 – 构建系统

作者: 2017 年 6 月 27 日2024 年 11 月 16 日暂无评论

第一篇博文中,我解释了如何生成可在 Python 解释器中使用的 torch.Tensor 对象。接下来,我将探讨 PyTorch 的构建系统。PyTorch 代码库包含各种组件:

  • 核心 Torch 库:TH、THC、THNN、THCUNN
  • 供应商库:CuDNN、NCCL
  • Python 扩展库
  • 附加的第三方库:NumPy、MKL、LAPACK

一个简单的 python setup.py install 命令是如何让你能够调用 import torch 并在代码中使用 PyTorch 库的呢?

本文的第一部分将从最终用户的角度解释构建过程。这将解释我们如何利用上述组件来构建库。文档的第二部分对 PyTorch 开发者很重要。它将记录通过仅构建你正在处理的代码子集来提高迭代速度的方法。

Setuptools 和 PyTorch 的 setup() 函数

Python 使用 Setuptools 来构建库。Setuptools 是核心 Python 库中原始 distutils 系统的扩展。Setuptools 的核心组件是 setup.py 文件,其中包含构建项目所需的所有信息。最重要的函数是 setup() 函数,它作为主入口点。让我们看看 PyTorch 中的那个:

setup(name="torch", version=version,
      description="Tensors and Dynamic neural networks in Python with strong GPU acceleration",
      ext_modules=extensions,
      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,
      },
      packages=packages,
      package_data={'torch': [
          'lib/*.so*', 'lib/*.dylib*',
          'lib/torch_shm_manager',
          'lib/*.h',
          'lib/include/TH/*.h', 'lib/include/TH/generic/*.h',
          'lib/include/THC/*.h', 'lib/include/THC/generic/*.h']},
      install_requires=['pyyaml'],
      )

该函数完全由关键字参数组成,这些参数有两个目的:

  • 元数据(例如名称、描述、版本)
  • 包的内容

我们关注第二点。让我们分解各个组件:

  • ext_modules:Python 模块要么是“纯”模块,只包含 Python 代码,要么是“扩展”模块,用 Python 实现的低级语言编写。这里我们列出了构建中的扩展模块,包括包含我们的 Python Tensor 的主 torch._C 库。
  • cmdclass:当从命令行使用 setup.py 脚本时,用户必须指定一个或多个“命令”,即执行特定操作的代码片段。例如,“install”命令构建并安装包。此映射将特定命令路由到 setup.py 中实现它们的函数。
  • packages:项目中的包列表。这些是“纯”的——即它们只包含 Python 代码。这些在 setup.py 中其他地方定义。
  • package_data:需要安装到包中的附加文件:在这种情况下,构建将生成的头文件和共享库必须包含在我们的安装中。
  • install_requires:为了构建 PyTorch,我们需要 pyyaml。Setuptools 将处理确保 pyyaml 可用,并在必要时下载和安装它。

我们将更详细地考虑这些组件,但现在,查看安装的最终产品——即 Setuptools 在构建代码后所做的事情——是很有启发性的。

site_packages

第三方包默认安装到与你的 Python 二进制文件关联的 lib//site_packages 目录中。例如,由于我使用的是 Miniconda 环境,我的 Python 二进制文件位于:

(p3) killeent@devgpu047:pytorch (master)$ which python
~/local/miniconda2/envs/p3/bin/python

因此包安装到:

/home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages

我安装了 PyTorch,让我们看看 site-packages 中的 torch 文件夹:

(p3) killeent@devgpu047:site-packages$ cd torch
(p3) killeent@devgpu047:torch$ ls
autograd  backends  _C.cpython-36m-x86_64-linux-gnu.so  cuda  distributed  _dl.cpython-36m-x86_64-linux-gnu.so  functional.py  __init__.py  legacy  lib  multiprocessing  nn  optim  __pycache__  serialization.py  _six.py  sparse  storage.py  _tensor_docs.py  tensor.py  _tensor_str.py  _thnn  _torch_docs.py  utils  _utils.py  version.py

请注意,我们期望在这里的一切都在这里:

  • 所有“纯”包都在这里 [待办,从 setup.py 打印包以解释]
  • 扩展库在这里——._C* 和 ._dl* 共享库
  • package_data 在这里:lib/ 的内容与我们在 setup 函数中描述的完全一致
(p3) killeent@devgpu047:torch$ ls lib/
include     libnccl.so.1  libTHC.so.1   libTHCUNN.so.1  libTHNN.so.1  libTH.so.1   THCUNN.h  torch_shm_manager libnccl.so  libshm.so     libTHCS.so.1  libTHD.so.1     libTHPP.so.1  libTHS.so.1  THNN.h

Python 解释器在导入期间会查看 site_packages。如果我们在 Python 代码中调用 import torch,它将在此处找到模块并初始化和导入它。你可以在这里阅读有关导入系统的更多信息。

构建各个部分

接下来,我们将从头到尾查看构建的各个组件。这将说明我们如何组合我们在引言中提到的所有代码。

后端 Torch 和供应商库

让我们看看 PyTorch 的 setup.py 中的 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”的命令——让我们看看它的 run() 方法:

def run(self):
        from tools.nnwrap import generate_wrappers as generate_nn_wrappers
        build_all_cmd = ['bash', 'torch/lib/build_all.sh']
        if WITH_CUDA:
            build_all_cmd += ['--with-cuda']
        if WITH_NCCL and not SYSTEM_NCCL:
            build_all_cmd += ['--with-nccl']
        if WITH_DISTRIBUTED:
            build_all_cmd += ['--with-distributed']
        if subprocess.call(build_all_cmd) != 0:
            sys.exit(1)
        generate_nn_wrappers()

在这里我们注意到 torch/lib/ 目录中有一个 shell 脚本 build_all.sh。该脚本可以通过是否启用了 CUDA、NCCL 库和 PyTorch 的分布式库来进行配置。

让我们看看 torch/lib

(p3) killeent@devgpu047:lib (master)$ ls
build_all.sh  libshm  nccl  README.md  TH  THC  THCS  THCUNN  THD  THNN  THPP  THS

在这里我们看到所有后端库的目录。THTHCTHNNTHCUNNncclgit 子树,它们与例如 github.com/torch 中的库同步。THSTHCSTHDTHPPlibshm 是 PyTorch 特有的库。所有库都包含 CMakeLists.txt——表明它们是用 CMake 构建的。

build_all.sh 本质上是一个脚本,它对所有这些库运行 CMake 配置步骤,然后运行 make install。让我们运行 ./build_all.sh,看看我们得到了什么:

(p3) killeent@devgpu047:lib (master)$ ./build_all.sh --with-cuda --with-nccl --with-distributed
[various CMake output logs]
(p3) killeent@devgpu047:lib (master)$ ls
build  build_all.sh  include  libnccl.so  libnccl.so.1  libshm  libshm.so  libTHC.so.1  libTHCS.so.1  libTHCUNN.so.1  libTHD.so.1  libTHNN.so.1  libTHPP.so.1  libTH.so.1  libTHS.so.1  nccl  README.md  TH  THC  THCS  THCUNN  THCUNN.h  THD  THNN  THNN.h  THPP  THS  tmp_install  torch_shm_manager

现在目录中有很多额外的东西:

  • 每个库的共享库文件
  • THNNTHCUNN 的头文件
  • buildtmp_install 目录
  • torch_shm_manager 可执行文件

让我们进一步探索。在 shell 脚本中,我们创建 build 目录和每个要构建的库的子目录:

# We create a build directory for the library, which will
# contain the cmake output. $1 is the library to be built
  mkdir -p build/$1
  cd build/$1

因此,例如 build/TH 包含 CMake 配置输出,包括用于构建 TH 的 Makefile,以及在此目录中运行 make install 的结果。

我们再看看 tmp_install

(p3) killeent@devgpu047:lib (master)$ ls tmp_install/
bin  include  lib  share

tmp_install 看起来像一个标准的安装目录,包含二进制文件、头文件和库文件。例如,tmp_install/include/TH 包含所有 TH 头文件,tmp_install/lib/ 包含 libTH.so.1 文件。

那么为什么要这个目录呢?它用于编译相互依赖的库。例如,THC 库依赖于 TH 库及其头文件。这在构建 shell 脚本中作为 cmake 命令的参数引用:

# install_dir is tmp_install
cmake ...
	-DTH_INCLUDE_PATH="$INSTALL_DIR/include" \
	-DTH_LIB_PATH="$INSTALL_DIR/lib" \

事实上,如果我们查看我们构建的 THC 库:

(p3) killeent@devgpu047:lib (master)$ ldd libTHC.so.1
	...
	libTH.so.1 => /home/killeent/github/pytorch/torch/lib/tmp_install/lib/./libTH.so.1 (0x00007f84478b7000)

build_all.sh 指定包含和库路径的方式有点混乱,但这代表了整体思想。最后,在脚本的末尾:

# If all the builds succeed we copy the libraries, headers,
# binaries to torch/lib
cp $INSTALL_DIR/lib/* .
cp THNN/generic/THNN.h .
cp THCUNN/generic/THCUNN.h .
cp -r $INSTALL_DIR/include .
cp $INSTALL_DIR/bin/* .

正如我们所看到的,最后,我们将所有内容复制到顶层 torch/lib 目录——解释了我们上面看到的内容。接下来我们将看到为什么这样做。

NN 包装器

简要地说,让我们谈谈 build_deps 命令的最后一部分:generate_nn_wrappers()。我们使用 PyTorch 的自定义 cwrap 工具绑定到后端库,我们在之前的文章中提到过。对于绑定 THTHC,我们手动为每个函数编写 YAML 声明。但是,由于 THNNTHCUNN 库相对简单,我们自动生成 cwrap 声明和生成的 C++ 代码。

我们将 THNN.hTHCUNN.h 头文件复制到 torch/lib 的原因是 generate_nn_wrappers() 代码期望这些文件位于此处。generate_nn_wrappers() 做了一些事情:

  1. 解析头文件,生成 cwrap YAML 声明并将它们写入输出 .cwrap 文件
  2. 使用适当的插件调用 cwrap 处理这些 .cwrap 文件,为每个文件生成源代码
  3. 第二次解析头文件,生成 THNN_generic.h——一个库,它接收 THPP 张量(PyTorch 的“通用”C++ 张量库),并根据张量的动态类型调用适当的 THNN/THCUNN 库函数

如果我们在运行 generate_nn_wrappers() 后查看 torch/csrc/nn,我们可以看到输出:

(p3) killeent@devgpu047:nn (master)$ ls
THCUNN.cpp  THCUNN.cwrap  THNN.cpp  THNN.cwrap  THNN_generic.cpp  THNN_generic.cwrap  THNN_generic.h  THNN_generic.inc.h

例如,代码生成 cwrap 如下:

[[
  name: FloatBatchNormalization_updateOutput
  return: void
  cname: THNN_FloatBatchNormalization_updateOutput
  arguments:
    - void* state
    - THFloatTensor* input
    - THFloatTensor* output
    - type: THFloatTensor*
      name: weight
      nullable: True
    - type: THFloatTensor*
      name: bias
      nullable: True
    - THFloatTensor* running_mean
    - THFloatTensor* running_var
    - THFloatTensor* save_mean
    - THFloatTensor* save_std
    - bool train
    - double momentum
    - double eps
]]

以及相应的 .cpp

extern "C" void THNN_FloatBatchNormalization_updateOutput(void*, THFloatTensor*, THFloatTensor*, THFloatTensor*, THFloatTensor*, THFloatTensor*, THFloatTensor*, THFloatTensor*, THFloatTensor*, bool, double, double);

PyObject * FloatBatchNormalization_updateOutput(PyObject *_unused, PyObject *args) {
	// argument checking, unpacking
	 PyThreadState *_save = NULL;
      try {
        Py_UNBLOCK_THREADS;
        THNN_FloatBatchNormalization_updateOutput(arg_state, arg_input, arg_output, arg_weight, arg_bias, arg_running_mean, arg_running_var, arg_save_mean, arg_save_std, arg_train, arg_momentum, arg_eps);
        Py_BLOCK_THREADS;
        Py_RETURN_NONE;
      } catch (...) {
        if (_save) {
          Py_BLOCK_THREADS;
        }
        throw;
      }

    ...
}

THPP 生成的代码中,函数如下所示:

void BatchNormalization_updateOutput(thpp::Tensor* input, thpp::Tensor* output, thpp::Tensor* weight, thpp::Tensor* bias, thpp::Tensor* running_mean, thpp::Tensor* running_var, thpp::Tensor* save_mean, thpp::Tensor* save_std, bool train, double momentum, double eps) {
	// Call appropriate THNN function based on tensor type, whether its on CUDA, etc.
}

我们稍后将更深入地了解这些源文件的用途。

“构建”纯 Python 模块

现在我们已经构建了后端库(“依赖项”),我们可以继续构建实际的 PyTorch 代码。接下来运行的 Setuptools 命令是 build_py,它用于构建我们库中的所有“纯”Python 模块。这些是传递给 setup.py 的“包”。

使用 Setuptools 的实用函数 find_packages() 查找包:

packages = find_packages(exclude=('tools.*',))
['torch', 'torch._thnn', 'torch.autograd', 'torch.backends', 'torch.cuda', 'torch.distributed', 'torch.legacy', 'torch.multiprocessing', 'torch.nn', 'torch.optim', 'torch.sparse', 'torch.utils', 'torch.autograd._functions', 'torch.backends.cudnn', 'torch.legacy.nn', 'torch.legacy.optim', 'torch.nn._functions', 'torch.nn.backends', 'torch.nn.modules', 'torch.nn.parallel', 'torch.nn.utils', 'torch.nn._functions.thnn', 'torch.utils.data', 'torch.utils.ffi', 'torch.utils.serialization', 'torch.utils.trainer', 'torch.utils.backcompat', 'torch.utils.trainer.plugins']

如我们所见,find_package 已经递归遍历了 torch 目录,找到了所有包含 __init__.py 文件的目录路径。

使用 Setuptools 构建时,该工具会在分发根目录(即 setup.py 文件所在的相同位置)中创建一个 build 目录。因为 PyTorch 由“纯”Python 模块和扩展模块组成,所以我们需要保留在执行构建时使用的操作系统和 Python 版本的信息。所以如果我们查看我的 build 目录,我们会看到:

(p3) killeent@devgpu047:pytorch (master)$ ls build
lib.linux-x86_64-3.6  temp.linux-x86_64-3.6

这表明我已经在 linux-x86-64 上使用 Python 3.6 构建了项目。lib 目录包含库文件,而 temp 目录包含构建过程中生成但在最终安装中不需要的文件。

由于“纯”Python 模块只是 Python 代码,不需要“编译”,因此 build_py 过程只是将文件从 find_packages 找到的位置复制到 build/ 中的等效位置。所以我们的构建输出充满了这样的行:

copying torch/autograd/_functions/blas.py -> build/lib.linux-x86_64-3.6/torch/autograd/_functions

我们之前还注意到,我们可以将文件和目录传递给主 setup() 函数的 package_data 关键字参数,并且 Setuptools 将处理将这些文件复制到安装位置。在 build_py 期间,这些文件被复制到 build/ 目录,所以我们也看到这样的行:

copying torch/lib/libTH.so.1 -> build/lib.linux-x86_64-3.6/torch/lib
...
copying torch/lib/include/THC/generic/THCTensor.h -> build/lib.linux-x86_64-3.6/torch/lib/include/THC/generic

构建扩展模块

最后,我们需要构建扩展模块,即使用 CPython 后端用 C++ 编写的 PyTorch 模块。这也构成了 setup.py 中大部分代码逻辑。我们重写的 build_ext 命令在扩展本身实际构建之前有一些特殊逻辑:

from tools.cwrap import cwrap
from tools.cwrap.plugins.THPPlugin import THPPlugin
from tools.cwrap.plugins.ArgcountSortPlugin import ArgcountSortPlugin
from tools.cwrap.plugins.AutoGPU import AutoGPU
from tools.cwrap.plugins.BoolOption import BoolOption
from tools.cwrap.plugins.KwargsPlugin import KwargsPlugin
from tools.cwrap.plugins.NullableArguments import NullableArguments
from tools.cwrap.plugins.CuDNNPlugin import CuDNNPlugin
from tools.cwrap.plugins.WrapDim import WrapDim
from tools.cwrap.plugins.AssertNDim import AssertNDim
from tools.cwrap.plugins.Broadcast import Broadcast
from tools.cwrap.plugins.ProcessorSpecificPlugin import ProcessorSpecificPlugin
        thp_plugin = THPPlugin()
        cwrap('torch/csrc/generic/TensorMethods.cwrap', plugins=[
            ProcessorSpecificPlugin(), BoolOption(), thp_plugin,
            AutoGPU(condition='IS_CUDA'), ArgcountSortPlugin(), KwargsPlugin(),
            AssertNDim(), WrapDim(), Broadcast()
        ])
        cwrap('torch/csrc/cudnn/cuDNN.cwrap', plugins=[
            CuDNNPlugin(), NullableArguments()
        ])

回想上面我记录了我们自动生成调用 THNN 等库的 C++ 代码。这就是我们绑定 THTHCCuDNN 的地方。我们采用 TensorMethods.cwrap 中的 YAML 声明,并使用它们生成包含在 PyTorch 的 C++ 生态系统中工作的实现的输出 C++ 源文件。例如,像 zero_ 这样的简单声明:

[[
  name: zero_
  cname: zero
  return: self
  arguments:
    - THTensor* self
]]

生成如下代码:

 PyObject * THPTensor_(zero_)(PyObject *self, PyObject *args, PyObject *kwargs) {
	...
	THTensor_(zero)(LIBRARY_STATE arg_self);
	...
}

在上一篇文章中,我们记录了这些函数如何绑定到特定的张量类型,所以这里就不再赘述了。对于构建过程来说,只要知道这些 C++ 文件是在扩展构建之前生成的就足够了,因为这些源文件在扩展编译期间被使用。

指定扩展

与纯模块不同,仅仅列出模块或包并期望 Setuptools 找到正确的文件是不够的;你必须指定扩展名、源文件以及任何编译/链接要求(包含目录、要链接的库等)。

setup.py 的大部分(撰写本文时大约 200 行代码)都用于指定如何构建这些扩展。在这里,我们在 build_all.sh 中做出的一些选择开始变得有意义。例如,我们看到我们的构建脚本指定了一个 tmp_install 目录,我们在其中安装了后端库。在我们的 setup.py 代码中,我们在添加到包含头文件的目录列表时引用此目录:

# tmp_install_path is torch/lib/tmp_install
include_dirs += [
    cwd,
    os.path.join(cwd, "torch", "csrc"),
    tmp_install_path + "/include",
    tmp_install_path + "/include/TH",
    tmp_install_path + "/include/THPP",
    tmp_install_path + "/include/THNN",

同样,我们在 build_all.sh 脚本的末尾将共享对象库复制到 torch/csrc。我们在 setup.py 代码中直接引用这些位置,以识别我们可能链接的库:

# lib_path is torch/lib
TH_LIB = os.path.join(lib_path, 'libTH.so.1')
THS_LIB = os.path.join(lib_path, 'libTHS.so.1')
THC_LIB = os.path.join(lib_path, 'libTHC.so.1')
THCS_LIB = os.path.join(lib_path, 'libTHCS.so.1')
THNN_LIB = os.path.join(lib_path, 'libTHNN.so.1')
# ...

让我们考虑如何构建主要的 torch._C 扩展模块:

C = Extension("torch._C",
              libraries=main_libraries,
              sources=main_sources,
              language='c++',
              extra_compile_args=main_compile_args + extra_compile_args,
              include_dirs=include_dirs,
              library_dirs=library_dirs,
              extra_link_args=extra_link_args + main_link_args + [make_relative_rpath('lib')],
              )
  • 主库是我们链接的所有库。这包括 shm(PyTorch 的共享内存管理库)以及 cudartcudnn 等系统库。请注意,这里没有列出 TH 库。
  • 主源文件是构成 PyTorch C++ 后端的 C++ 文件。
  • 编译参数是配置编译的各种标志。例如,我们可能希望在调试模式下编译时添加调试标志。
  • 包含目录是所有包含头文件的目录的路径。这也是 build_all.sh 脚本重要的另一个例子——例如,我们在 torch/lib/tmp_install/include/TH 中查找 TH 头文件——这是我们通过 CMake 配置指定的安装位置。
  • 库目录是在链接时搜索共享库的目录。例如,我们包括 torch/lib——我们在 build_all.sh 结束时将 .so 文件复制到的位置,以及 CUDA 和 CuDNN 目录的路径。
  • 链接参数用于将目标文件链接在一起以创建扩展。在 PyTorch 中,这包括更正常的选项,例如决定静态链接 libstdc++。但是,有一个关键组件:这就是我们链接后端 TH 库的地方。请注意,我们有如下行:
# The explicit paths to .so files we described above
main_link_args = [TH_LIB, THS_LIB, THPP_LIB, THNN_LIB]

你可能想知道我们为什么要这样做,而不是将这些库添加到我们传递给 libraries 关键字参数的列表中。毕竟,这是一个要链接的库列表。问题是 Lua Torch 安装通常会设置 LD_LIBRARY_PATH 变量,因此我们可能会错误地链接到为 Lua Torch 构建的 TH 库,而不是我们本地构建的库。这将带来问题,因为代码可能已过时,而且 Lua Torch 的 TH 有各种配置选项,可能与 PyTorch 不兼容。

因此,我们直接向链接器手动指定我们生成的共享库的路径。

PyTorch 还需要其他扩展,它们的构建方式类似。Setuptools 库调用 C++ 编译器和链接器来构建所有这些扩展。如果构建成功,我们就成功构建了 PyTorch 库,然后可以继续安装。

安装

构建完成后,安装非常简单。我们只需将所有内容从 build/lib.linux-x86_64-3.6 目录复制到相应的安装目录。回想一下,我们上面指出此目录是与我们的 Python 二进制文件关联的 site_packages 目录。因此,我们看到如下行:

running install_lib
creating /home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/torch
copying build/lib.linux-x86_64-3.6/torch/_C.cpython-36m-x86_64-linux-gnu.so -> /home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/torch
copying build/lib.linux-x86_64-3.6/torch/_dl.cpython-36m-x86_64-linux-gnu.so -> /home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/torch
creating /home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/torch/_thnn
copying build/lib.linux-x86_64-3.6/torch/_thnn/_THNN.cpython-36m-x86_64-linux-gnu.so -> /home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/torch/_thnn
copying build/lib.linux-x86_64-3.6/torch/_thnn/_THCUNN.cpython-36m-x86_64-linux-gnu.so -> /home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/torch/_thnn

最后,让我们启动 Python 解释器。当 Python 解释器执行导入语句时,它会沿着搜索路径搜索 Python 代码和扩展模块。路径的默认值在构建解释器时配置到 Python 二进制文件中。

# note we are now in my home directory
(p3) killeent@devgpu047:~$ python
Python 3.6.1 |Continuum Analytics, Inc.| (default, Mar 22 2017, 19:54:23)
[GCC 4.4.7 20120313 (Red Hat 4.4.7-1)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import sys
>>> sys.path
['', '/home/killeent/local/miniconda2/envs/p3/lib/python36.zip', '/home/killeent/local/miniconda2/envs/p3/lib/python3.6', '/home/killeent/local/miniconda2/envs/p3/lib/python3.6/lib-dynload', '/home/killeent/.local/lib/python3.6/site-packages', '/home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages', '/home/killeent/github/pytorch', '/home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/setuptools-27.2.0-py3.6.egg']

如我们所见,我们复制 PyTorch 安装的 site-packages 目录是搜索路径的一部分。现在让我们加载 torch 模块并查看其位置:

>>> import torch
>>> import inspect
>>> inspect.getfile(torch)
'/home/killeent/local/miniconda2/envs/p3/lib/python3.6/site-packages/torch/__init__.py'

如我们所见,我们已按预期从 site_packages 加载了模块——并且我们的构建和安装成功!

注意:Python 在 sys.path 前面加上空字符串,表示当前工作目录——使其成为我们搜索模块的第一个位置。因此,如果我们在 pytorch 目录中运行 Python,我们可能会意外加载本地版本的 PyTorch,而不是我们安装的版本。这是一件需要注意的事情。

附录 – 开发者效率、第三方库、我没有涵盖的内容

PyTorch 的整个安装循环可能非常耗时。在我的开发服务器上,从源代码安装大约需要 5 分钟。通常,在开发 PyTorch 时,我们只想处理整个项目的一个子集,并仅重新构建该子集以测试更改。幸运的是,我们的构建系统支持这一点。

Setuptools 开发模式

支持此功能的主要工具是 Setuptools 的 develop 命令。文档指出:

此命令允许你部署项目的源代码以在其中一个或多个“暂存区”中使用,以便导入。此部署方式使得项目源代码的更改立即可在暂存区中使用,而无需在每次更改后运行构建或安装步骤。

但是它是如何工作的呢?假设我们在 PyTorch 目录中运行 python setup.py build developbuild 命令运行,构建我们的依赖项(THTHPP 等)和扩展库。然而,如果我们查看 site-packages

(p3) killeent@devgpu047:site-packages$ ls -la torch*
-rw-r--r--. 1 killeent users 31 Jun 27 08:02 torch.egg-link

查看 torch.egg-link 文件的内容,它只是引用了 PyTorch 目录:

(p3) killeent@devgpu047:site-packages$ cat torch.egg-link
/home/killeent/github/pytorch

如果我们回到 PyTorch 目录,我们会看到有一个新的目录 torch.egg-info

(p3) killeent@devgpu047:pytorch (master)$ ls -la torch.egg-info/
total 28
drwxr-xr-x.  2 killeent users  4096 Jun 27 08:09 .
drwxr-xr-x. 10 killeent users  4096 Jun 27 08:01 ..
-rw-r--r--.  1 killeent users     1 Jun 27 08:01 dependency_links.txt
-rw-r--r--.  1 killeent users   255 Jun 27 08:01 PKG-INFO
-rw-r--r--.  1 killeent users     7 Jun 27 08:01 requires.txt
-rw-r--r--.  1 killeent users 16080 Jun 27 08:01 SOURCES.txt
-rw-r--r--.  1 killeent users    12 Jun 27 08:01 top_level.txt

此文件包含有关 PyTorch 项目的元数据。例如,requirements.txt 列出了设置 PyTorch 的所有依赖项:

(p3) killeent@devgpu047:pytorch (master)$ cat torch.egg-info/requires.txt
pyyaml

不深入细节,develop 允许我们本质上将 PyTorch 仓库本身视为位于 site-packages 中,因此我们可以导入模块并使其正常工作:

(p3) killeent@devgpu047:~$ python
Python 3.6.1 |Continuum Analytics, Inc.| (default, Mar 22 2017, 19:54:23)
[GCC 4.4.7 20120313 (Red Hat 4.4.7-1)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__file__
'/home/killeent/github/pytorch/torch/__init__.py'

因此,以下后果成立:

  • 如果我们更改 Python 源文件,更改会自动被检测到,并且我们无需运行任何命令即可让 Python 解释器看到此更改。
  • 如果我们更改了其中一个扩展库中的 C++ 源文件,我们可以重新运行 develop 命令,它将重新构建扩展。

因此,我们可以无缝开发 PyTorch 代码库,并以简单的方式测试我们的更改。

处理依赖库

如果我们在处理依赖项(例如 THTHPP 等),我们可以通过直接运行 build_deps 命令更快地重新构建更改。这将自动调用 build_all.sh 来重新构建我们的库,并适当复制生成的库。如果我们在使用 Setuptools develop 模式,我们将使用 PyTorch 目录中构建的本地扩展库。因为我们在编译扩展库时指定了共享库的路径,所以更改将被检测到:

# we are using the local extension
(p3) killeent@devgpu047:~$ python
Python 3.6.1 |Continuum Analytics, Inc.| (default, Mar 22 2017, 19:54:23)
[GCC 4.4.7 20120313 (Red Hat 4.4.7-1)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch._C.__file__
'/home/killeent/github/pytorch/torch/_C.cpython-36m-x86_64-linux-gnu.so'

# it references the local shared object library we just re-built
(p3) killeent@devgpu047:~$ ldd /home/killeent/github/pytorch/torch/_C.cpython-36m-x86_64-linux-gnu.so
# ...
libTH.so.1 => /home/killeent/github/pytorch/torch/lib/libTH.so.1 (0x00007f543d0e2000)
# ...

因此,我们可以在不进行完全重新构建的情况下测试这里的任何更改。

第三方库

PyTorch 依赖一些第三方库。使用这些库的常见机制是通过 Anaconda 安装它们,然后链接它们。例如,我们可以使用以下方式将 mkl 库与 PyTorch 一起使用:

# installed to miniconda2/envs/p3/lib/libmkl_intel_lp64.so
conda install mkl

然后,只要我们的 $CMAKE_PREFIX_PATH 上有此 lib 目录的路径,它在编译时就能成功找到此库。

# in the site-packages dir
(p3) killeent@devgpu047:torch$ ldd _C.cpython-36m-x86_64-linux-gnu.so
# ...
libmkl_intel_lp64.so => /home/killeent/local/miniconda2/envs/p3/lib/libmkl_intel_lp64.so (0x00007f3450bba000)
# ...

未涵盖但同样相关的内容

  • 如何使用 ccache 来加速构建时间
  • PyTorch 的顶级 __init__.py 文件如何处理初始模块导入以及将所有各种模块和扩展库整合在一起
  • CMake 构建系统,以及后端库如何使用 CMake 进行配置和构建