pytorch下import torch_PyTorch（二）

依然采取自顶向下的原则剖析，借助PyTorch的python接口。我们知道使用PyTorch第一步都是

import torch

于是阅读torch/__init__.py，发现需要加载torch._C这个库，但是需要以（RTLD_GLOBAL|RTLD_LAZY）这个模式动态加载，于是先将动态加载模式设置到（RTLD_GLOBAL|RTLD_LAZY）之后加载torch._C然后再恢复动态加载模式，

old_flags=sys.getdlopenflags()
sys.setdlopenflags(_dl_flags.RTDL_GLOBAL | _dl_flags.RTLD_LAZY)
from torch._C import *
__all__ += [name for name in dir(_C)
            if name[0] != '_' and
            not name.endswith('Base')]
sys.setdlopenflags(old_flags)

将torch._C中（不包括_开头和Base结尾）的属性导出到当前域。

__init__.py除了import torch._C，还import了同目录下其他module，以及同目录下的package。首先看torch._C导入时做了什么， torch._C的源文件只有torch/csrc/stub.cpp，链接库为shm和torch_python，stub.cpp中仅仅是初始化模块，

extern PyObject* initModule();
PyMODINIT_FUNC PyInit__C()   // 在python脚本中，import _C 时调用
{
  return initModule();
}

根据python3扩展库的规则可知，import torch._C ，调用PyInit__C函数（调用名为PyInit_<package>的函数），这个函数内部调用initModule，也就是说，具体的模块定义由initModule实现。看到extern知道initModule方法定义在外部，所以只能从shm和torch_python对应的源文件中寻找方法定义。

shm库实现Domain Socket通信获得共享内存的句柄，解决多进程的内存分配问题，查看torch/CMakeLists.txt，发现生成shm相关语句为，

set(LIBSHM_SUBDIR libshm)
set(LIBSHM_SRCDIR ${LIBSHM_SRC_DIR}/lib/${LIBSHM_SUBDIR})
add_subdirectory(${LIBSHM_SRCDIR})

从上面语句得知shm库的源码位于torch/lib/libshm目录下，这个跟torch._C模块定义没有关系，暂且不细展开，继续查看torch_python的源码以寻求initModule方法定义。在torch/CMakeLists.txt中发现

add_library(torch_python SHARED ${TORCH_PYTHON_SRCS})

TORCH_PYTHON_SRCS是一个列表，存储了torch_python库的源文件，生成torch_python库所需要的源文件以及依赖库直接查看torch/CMakeLists.txt，这里不再展开一一说明。

initModule方法定义在torch/csrc/Module.cpp，

#ifdef USE_CUDA
namespace torch {
     namespace cuda {
    
void initModule(PyObject* module);       // 模块中有关cuda部分的初始化函数声明
}}
#endif

static std::vector<PyMethodDef> methods;

PyObject* module;
PyObject* initModule() {
                     // 声明并定义模块初始化函数
  // 向methods中添加方法定义
  THPUtils_addPyMethodDefs(methods, TorchMethods);
  THPUtils_addPyMethodDefs(methods, DataLoaderMethods);
  ...
  // 真正的扩展模块定义
  static struct PyModuleDef torchmodule = {
    
    PyModuleDef_HEAD_INIT,
    "torch._C",                          // 扩展模块名
    nullptr,                           
    -1,
    methods.data()                       // 模块中的方法定义
  };
  ASSERT_TRUE(module = PyModule_Create(&torchmodule)); // 创建模块并确保创建成功
  // 对模块进行各种初始化
#ifdef USE_CUDA
  torch::cuda::initModule(module);       // 执行cuda相关的初始化
#endif
  ...
  // 定义模块的属性设置函数，setter
  // 属性名为name，值为v，incref表示是否对值对象增加引用计数
  // 设置成功返回1，否则返回0
  auto set_module_attr = [&](const char* name, PyObject* v, bool incref = true) {
    
    if(incref) {
    
      Py_INCREF(v);
    }
    return PyModule_AddObject(module, name, v) == 0;
  }
  // 设置模块属性
  ...
  ASSERT_TRUE(set_module_attr("has_cudnn", has_cudnn));
  // 向模块添加方法
  auto py_module = py::reinterpret_borrow<py::module>(module);
  py_module.def("_demangle", &c10::demangle);
  py_module.def("_log_api_usage_once", &LogAPIUsageOnceFromPython);
  ...    // 设置模块其他属性
  ASSERT_TRUE(set_module_attr("default_generator", (PyObject*)THPDefaultGenerator, false));
  torch::nn::init__THNN(module);  // 增加 _THNN 属性
#ifdef USE_CUDA
  torch::nn::init_THCUDD(module);
#endif
  return module;
  ...
}

从上面的代码中可见，定义并生成名为torch._C的模块，然后对这个模块设置attr，添加方法，添加子模块等，

• 使用 THPUtils_addPyMethodDefs 向torch._C 添加模块方法。包括

# TorchMethods 
_initExtension
_autograd_init
...
# DataLoaderMethods 
_set_worker_signal_handlers
_set_worker_pids
...
# torch::autograd::python_functions(), torch/csrc/autograd/init.cpp
set_grad_enabled
is_grad_enabled
set_anomaly_enabled
is_anomaly_enabled
# torch::multiprocessing::python_functions(), torch/csrc/multiprocessing/init.cpp
_multiprocessing_init
# torch::distributed::c10d::python_functions()  同上类似
...
# THCPModule_method(), torch/csrc/cuda/Module.cpp
_cuda_init
_cuda_setDevice
...
_nccl_version
...
# THCUDNN_method()
_cudnn_version
# THDPModule_methods(), torch/csrc/distributed/Module.cpp
_dist_init_extension
_dist_init_process_group
...

• 生成模块torch._C 后再向其添加如下字段：

a. 向torch._C添加类型_PtrWrapper，Generator，FatalError，Size，dtype，iinfo，layout，memory_format，device，_LegacyVariableBase，_TensorBase，_VariableFunctions，_FunctionBase，_EngineBase，JITException，IODescriptor，_THNN，_THCUNN。

torch._C._TensorBase这个类型具有属性

_cdata
_version
grad_fn
_grad_fn
is_leaf
data
_grad
grad
...
device
ndim

并且具有以下方法

# variable_methods, torch/csrc/autograd/generated/python_variable_methods.cpp
__add__
__radd__
...
apply_
byte
char
contiguous
...
where
zero_
# extra_method
_make_subclass

类型torch._C._FunctionBase， 这个类型具有方法和属性为

# method
apply
_do_forward
_do_backward
_register_hook_dict
register_hook
# property
saved_tensors
saved_variables
...
requires_grad
metadata

不难知道_TensorBase是