作者:Trevor Killeen

在第一篇文章中,我解释了如何生成一个您可以在 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() 函数。

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/<version>/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

在这里,我们看到所有后端库的目录。THTHCTHNNTHCUNNnccl 是与例如 github.com/torch 中的库同步的git subtreeTHSTHCSTHDTHPPlibshm 是 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

现在目录中多了几样东西:

  • 每个库的共享库文件
  • THNN 和 THCUNN 的头文件
  • build 和 tmp_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 Tensor(PyTorch 的“通用”C++ Tensor 库),并根据 Tensor 的动态类型调用相应的 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 的“packages”。

包是使用 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 声明,并使用它们生成 C++ 输出源文件,这些文件包含可在 PyTorch 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);
	...
}

在上一篇文章中,我们记录了这些函数如何与特定的 Tensor 类型关联,所以在此不再赘述。对于构建过程来说,只需知道这些 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')],
              )
  • main libraries 是我们链接的所有库。这包括像 shm(PyTorch 的共享内存管理库)以及像 cudartcudnn 这样的系统库。请注意,TH 库没有在此列出。
  • main sources 是构成 PyTorch C++ 后端的 C++ 文件。
  • compile args 是配置编译的各种标志。例如,在调试模式下编译时,我们可能希望添加调试标志。
  • include dirs 是所有包含头文件的目录路径。这也是 build_all.sh 脚本很重要的另一个例子 - 例如,我们在 torch/lib/tmp_install/include/TH 中查找 TH 头文件 - 这是我们用 CMake 配置指定的安装位置。
  • library dirs 是在链接时搜索共享库的目录。例如,我们包含 torch/lib - 这是我们在 build_all.sh 末尾将 .so 文件复制到的位置,但也包含 CUDA 和 CuDNN 目录的路径。
  • link arguments 用于将目标文件链接在一起以创建扩展模块。在 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 解释器执行 import 语句时,它会在搜索路径中查找 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 的 Develop 模式

支持此功能的主要工具是 Setuptools 的 develop 命令。文档中写道:

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

但它是如何工作的呢?假设我们在 PyTorch 目录中运行 python setup.py build developbuild 命令会运行,构建我们的依赖项(TH、THPP 等)和扩展库。然而,如果我们查看 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 安装它们,然后进行链接。例如,我们可以通过执行以下操作在 PyTorch 中使用 mkl 库:

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

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

# 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 配置和构建