小白学习pytorch源码（二）：setup.py最详细解读

pytorch setup.py最全解析

从pytorch源码整体学习的角度看，有两个文件最为关键，分别为pytorch源码总目录下的setup.py和torch包中的_init_.py。其中_init_.py已经在我的上一篇文章中详细讲解过，这里不多赘述，感兴趣读者可以在下文链接找到，本期的阅读源码主要聚焦于setup.py这一文件。

setup.py与setuptools

当学习一个文件时，我们首先要明晰该文件的作用是什么。对于setup.py来说，它的作用是将pytorch文件进行安装（包括编译c/c++文件，检查dll文件和模组文件是否完全等等），并将该项目安装到当前环境python的‘site-packages’目录下，使其可以像导入标准库一样导入。要完成该功能，pytorch开发人员使用了setuptools工具，所以其实setup.py是按照setuptools的规定格式编写的。因此想要彻底理解setup.py的全部代码含义，首先要完全理解setuptools的使用方法。在后文讲解到pytorch相关代码时会引用到部分setuptools的知识，这里作者将setuptools的官方文档贴出，如果想详细了解setuptools工具，可以自由查阅。

setup.py最详细解读

setup.py 环境检查

关于setup.py的解读，我会从头到尾的梳理一遍代码的作用，以便读者阅览。一开始，该文件在导入了相关的setuptools和一些自定义工具文件后，对所在主机的环境进行检查，并依照传入参数对安装方式进行设置，同时检查文件是否完全，确定所需文件的位置为之后的安装做准备。
在这之后，我们就可以把目光转向setup.py的模拟程序入口，即_main_()，看看如何进行安装。_main_()函数本身函数并不长，在这里贴出：

if __name__ == '__main__':
    # Parse the command line and check the arguments
    # before we proceed with building deps and setup
    dist = Distribution()
    try:
        dist.parse_command_line()
    except setuptools.distutils.errors.DistutilsArgError as e:
        print(e)
        sys.exit(1)

    mirror_files_into_torchgen()
    if RUN_BUILD_DEPS:
        build_deps()

    extensions, cmdclass, packages, entry_points, extra_install_requires = configure_extension_build()

    install_requires += extra_install_requires

    # Read in README.md for our long_description
    with open(os.path.join(cwd, "README.md"), encoding="utf-8") as f:
        long_description = f.read()

    version_range_max = max(sys.version_info[1], 9) + 1
    setup(
    	##这里省略，原因后文讲解
    )
    if EMIT_BUILD_WARNING:
        print_box(build_update_message)

首先，文件实例化了一个Distribution()类变量。Distribution()类是定义在setuptools中的一个类，其作用就是检查主机环境是否含有项目所要求的所有的依赖，若没有则自动安装。
而下文的mirror_files_into_torchgen()则是检查项目的一些临时文件是否已经被创建，若没有则创建目录和文件，保证本机环境和torch需要的安装环境相同。
build_dep()函数也是检查工作的一部分，其内部包含检查分模块，检查python包依赖以及python文件位置等等。
值得注意的是，这些检查行为都发生在真正的setup操作之前，目的在于保证安装的正确进行，是安装pytorch必不可少的操作。

接下来，我们要详细了解一下configure_extension_build()，这一函数内说明了c/c++文件在配置阶段是如何与torch项目相关联的。在该函数内部，我们要对这一段代码加以重视：

   ################################################################################
    # Declare extensions and package
    ################################################################################

    extensions = []
    packages = find_packages(exclude=('tools', 'tools.*'))
    C = Extension("torch._C",
                  libraries=main_libraries,
                  sources=main_sources,
                  language='c',
                  extra_compile_args=main_compile_args + extra_compile_args,
                  include_dirs=[],
                  library_dirs=library_dirs,
                  extra_link_args=extra_link_args + main_link_args + make_relative_rpath_args('lib'))
    C_flatbuffer = Extension("torch._C_flatbuffer",
                             libraries=main_libraries,
                             sources=["torch/csrc/stub_with_flatbuffer.c"],
                             language='c',
                             extra_compile_args=main_compile_args + extra_compile_args,
                             include_dirs=[],
                             library_dirs=library_dirs,
                             extra_link_args=extra_link_args + main_link_args + make_relative_rpath_args('lib'))
    extensions.append(C)
    extensions.append(C_flatbuffer)

我们可以清楚地看到，pytorch项目是在此处将c/c++语言作为Extension纳入自己的项目，并预备在后文的setup()进行编译的。注意此处的Extension类也为setuptools定义好的一个类，具体的作用将实例化时传入的各种参数包装好，留待setup()时一并使用。

到目前为止，setup.py一直到setup()之前的函数行为都已经讲解完了，如此多的准备和检查也是pytorch项目如此稳定和成功的基础。下面我们一起学习一下setup()函数，尤其是c/c++语言在其中是如何被编译的。

setup.py setup()函数

setup()函数无疑是setup.py()的核心，其本身也是由setuptools实现好的函数，pytorch项目将各种该函数需要的参数设计好，再传入setup()函数中，即可方便快捷的进行项目部署。setup()函数部分代码展示如下：

setup(
        name=package_name,
        version=version,
        description=("Tensors and Dynamic neural networks in "
                     "Python with strong GPU acceleration"),
        long_description=long_description,
        long_description_content_type="text/markdown",
        ext_modules=extensions,
        cmdclass=cmdclass,
        packages=packages,
        entry_points=entry_points,
        install_requires=install_requires,
        package_data={
            'torch': [
                'py.typed',
                'bin/*',
                'test/*',
                '_C/*.pyi',
                '_C_flatbuffer/*.pyi',
				##过多所以只展示部分
                'utils/model_dump/skeleton.html',
                'utils/model_dump/code.js',
                'utils/model_dump/*.mjs',
            ],
            'torchgen': [
                'packaged/ATen/*',
                'packaged/ATen/native/*',
                'packaged/ATen/templates/*',
            ],
            'caffe2': [
                'python/serialized_test/data/operator_test/*.zip',
            ],
        },
        url='https://pytorch.org/',
        download_url='https://github.com/pytorch/pytorch/tags',
        author='PyTorch Team',
        author_email='[email protected]',
        python_requires='>={}'.format(python_min_version_str),
        # PyPI package information.
        classifiers=[
            'Development Status :: 5 - Production/Stable',
            'Intended Audience :: Developers',
            'Intended Audience :: Education',
            ##过多所以只展示部分
            'Programming Language :: C++',
            'Programming Language :: Python :: 3',
        ] + ['Programming Language :: Python :: 3.{}'.format(i) for i in range(python_min_version[1], version_range_max)],
        license='BSD-3',
        keywords='pytorch machine learning',
    )

我们可以清楚地看到setup()函数参数包括了项目信息、项目包需求、扩展信息等等，感兴趣的话读者可以自行展开浏览全部，在这里我请大家重点关注以下一行代码：

ext_modules=extensions

很明显，这行代码将刚才我们提到的c/c++扩展包放入了安装过程中，那么具体是怎样安装的呢？
我们回到之前提到的configure_extension_build()函数中与c/c++相关的片段，注意到在声明Extension类时，其中提到过sources参数，而其所传递的main_sources变量在上文被定义过，即：

main_sources = ["torch/csrc/stub.c"]

显然，如何编译c++语言与这一文件由很强的相关性，我们找到这一文件，其内部关键代码如下：

PyMODINIT_FUNC PyInit__C(void)
{
  return initModule();
}

再进一步探寻initModule()方法的来源，我们发现其被定义在/csrc/Module.cpp文件中，而阅览这个函数后，发现其应为整个pytorch项目初始化的关键函数，截取部分代码示例如下：

PyObject* initModule() {
  HANDLE_TH_ERRORS

  c10::initLogging();

  at::internal::lazy_init_num_threads();

  C10_LOG_API_USAGE_ONCE("torch.python.import");

// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
#define ASSERT_TRUE(cmd) \
  if (!(cmd))            \
  return nullptr

  THPUtils_addPyMethodDefs(methods, TorchMethods);
  THPUtils_addPyMethodDefs(methods, DataLoaderMethods);
  THPUtils_addPyMethodDefs(methods, torch::autograd::python_functions());
  THPUtils_addPyMethodDefs(methods, torch::multiprocessing::python_functions());
#ifdef USE_CUDA
  THPUtils_addPyMethodDefs(methods, THCPModule_methods());
#endif
#if defined(USE_DISTRIBUTED) && defined(USE_C10D)
  THPUtils_addPyMethodDefs(
      methods, torch::distributed::c10d::python_functions());
#ifndef _WIN32
  THPUtils_addPyMethodDefs(
      methods, torch::distributed::rpc::python_functions());
  THPUtils_addPyMethodDefs(
      methods, torch::distributed::autograd::python_functions());
  THPUtils_addPyMethodDefs(
      methods, torch::distributed::rpc::testing::python_functions());
#endif
#endif

  static struct PyModuleDef torchmodule = {
      PyModuleDef_HEAD_INIT, "torch._C", nullptr, -1, methods.data()};
  ASSERT_TRUE(module = PyModule_Create(&torchmodule));
  ASSERT_TRUE(THPGenerator_init(module));
  ASSERT_TRUE(THPException_init(module));
  THPSize_init(module);
  THPDtype_init(module);
  THPDTypeInfo_init(module);
  THPLayout_init(module);
  THPMemoryFormat_init(module);
  THPQScheme_init(module);
  THPDevice_init(module);
  THPStream_init(module);
  ASSERT_TRUE(THPVariable_initModule(module));
  ASSERT_TRUE(THPFunction_initModule(module));
  ASSERT_TRUE(THPEngine_initModule(module));

  ##下文过多不做赘述
  }

可以看到，其不但包括python函数的初始化，还包括一些项目设备、内存格式、设备信息等的初始化，还包括我们上一篇文章提到的THPVariable_initModule()函数，即使用c++实现torch包函数的初始化函数，代码后文还包括一些pybind11绑定相关代码等等。
到这一步时，我们已经通过对setup.py的研读，基本清楚了pytorch是如何安装并成为我们可以导入的python包的。
在之后的学习中，我会分模块学习pytorch的源代码以及底层实现，并发表出学习笔记，如果感兴趣的话请关注我，给我继续创作的动力，谢谢！

系列其他文章链接如下，持续更新中~
小白学习pytorch源码（一）：torch包函数如何实现？揭秘__init__.py
小白学习pytorch源码（二）：setup.py最详细解读