背景

在使用PyTorch深度學(xué)習(xí)框架的時候，不管是訓(xùn)練還是測試，代碼中引入PyTorch的第一句總是：

import torch

在Gemfield前述專欄文章里，我們已經(jīng)得知，torch/csrc/stub.cpp鏈接libshm.so、libtorch_python.so、libcaffe2_gpu.so生成了_C.cpython-37m-x86_64-linux-gnu.so庫，而像前述方式import torch的時候，按照python規(guī)范，會找到torch package目錄下的__init__.py，在這個文件中進(jìn)一步會調(diào)用：

from torch._C import *

其中torch._C就是_C.cpython-37m-x86_64-linux-gnu.so。因?yàn)?以Python3為例)按照Python規(guī)范，由于默認(rèn)的引擎都是CPython，而CPython的C/C++擴(kuò)展是一個共享庫，并且這個共享庫安裝在PYTHONPATH目錄下，并且文件名(不包含后綴)要和module的名字一樣，并且這個共享庫中要實(shí)現(xiàn)PyInit_modulename符號來作為import時候的邏輯入口。

對于PyTorch來說這個modulename 是_C，因此我們可以揣測，在torch/csrc/stub.cpp中一定實(shí)現(xiàn)了PyInit_C這個函數(shù)。是的，PyTorch就是這么做的，torch/csrc/stub.cpp中的代碼就是下面這樣：

#include

extern PyObject* initModule();

PyMODINIT_FUNC PyInit__C()

{

return initModule();

}

本文將從initModule函數(shù)展開，全面闡述PyTorch框架的初始化工作。initModule就是PyTorch初始化時候的第一層調(diào)用棧了，因?yàn)樗械某跏蓟ぷ鞫际窃谶@個函數(shù)內(nèi)完成的，內(nèi)容比較多，gemfield將其劃分為7部分：

1，torch._C的誕生：

這一步就是產(chǎn)生torch._C類，并在這個python類上面注冊眾多函數(shù)：

PyObject* initModule() {

//openmp的設(shè)置

THInferNumThreads();

THPUtils_addPyMethodDefs(methods, TorchMethods);

THPUtils_addPyMethodDefs(methods, DataLoaderMethods);

THPUtils_addPyMethodDefs(methods, torch::autograd::python_functions());

THPUtils_addPyMethodDefs(methods, torch::multiprocessing::python_functions());

THPUtils_addPyMethodDefs(methods, THCPModule_methods());

THPUtils_addPyMethodDefs(methods, THCUDNN_methods());

THPUtils_addPyMethodDefs(methods, THDPModule_methods());

THPUtils_addPyMethodDefs(methods, torch::distributed::c10d::python_functions());

module = Py_InitModule("torch._C", methods.data());

......

}

其中TorchMethods注冊了29個方法，都是THPModule_前綴的函數(shù)；DataLoaderMethods注冊了4個方法，都是THPModule_前綴的函數(shù)；torch::autograd::python_functions注冊了4個方法；torch::multiprocessing::python_functions注冊了1個方法；THCPModule_methods注冊了37個CUDA相關(guān)的函數(shù)，前綴都是THCPModule_；THCUDNN_methods注冊了1個方法；THDPModule_methods注冊了28個方法；torch::distributed::c10d::python_functions注冊了1個方法。

總而言之，在這一小步，我們達(dá)到了這樣一個里程碑，torch._C符號誕生，并且向torch._C注冊了一百余個函數(shù)，涉及torch、dataloader、autograd、multiprocess、cuda、cudnn、distribute、c10d方面。

2，一些關(guān)鍵類型

以下代碼先后初始化了torch._C._PtrWrapper、torch._C.Generator(含5個方法)、FatalError、torch.Size、torch.dtype、torch.iinfo、torch.layout、torch.device：

PyObject* initModule() {

......

THPWrapper_init(module);

THPGenerator_init(module);

THPException_init(module);

THPSize_init(module);

THPDtype_init(module);

THPDTypeInfo_init(module);

THPLayout_init(module);

THPDevice_init(module);

THPVariable_initModule(module);

THPFunction_initModule(module);

THPEngine_initModule(module);

......

}

3，torch._C._TensorBase的誕生

Gemfield將以下三個初始化函數(shù)歸為這一小節(jié)：

PyObject* initModule() {

......

THPVariable_initModule(module);

THPFunction_initModule(module);

THPEngine_initModule(module);

......

}

為什么呢？因?yàn)榈匚惶@赫了。

THPVariable_initModule(module) 創(chuàng)建了torch._C._TensorBase，這是一切Tensor的基類，在Gemfield的其它專欄文章里將單獨(dú)解釋；

THPFunction_initModule(module)創(chuàng)建了torch._C._FunctionBase，在torch/autograd/function.py中，以下兩個類以torch._C._FunctionBase為基類：

class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin))

class BackwardCFunction(_C._FunctionBase, _ContextMethodMixin, _HookMixin)

這個Function繼承體系就構(gòu)成了DAG的基礎(chǔ)。

THPEngine_initModule(module)創(chuàng)建了torch._C._EngineBase，_EngineBase這個類負(fù)責(zé)動態(tài)圖執(zhí)行之前的preprocess，_EngineBase會將torch.autograd的backward之類的請求預(yù)處理后送給真正的Engine去執(zhí)行。

4，pybind11綁定

這一小節(jié)的初始化內(nèi)容都是和pybind11相關(guān)的：

PyObject* initModule() {

......

// NOTE: We need to be able to access OperatorExportTypes from ONNX for use in

// the export side of JIT, so this ONNX init needs to appear before the JIT

// init.

torch::onnx::initONNXBindings(module);

torch::jit::initJITBindings(module);

torch::autograd::initNNFunctions(module);

torch::autograd::init_legacy_variable(module);

torch::python::init_bindings(module);

torch::cuda::initModule(module);

......

}

initONNXBindings是ONNX的python binding：torch._C._onnx.TensorProtoDataType和torch._C._onnx.OperatorExportTypes：

>>> dir(torch._C._onnx.TensorProtoDataType)

['BOOL', 'COMPLEX128', 'COMPLEX64', 'DOUBLE', 'FLOAT', 'FLOAT16', 'INT16', 'INT32', 'INT64', 'INT8', 'STRING', 'UINT16', 'UINT32', 'UINT64', 'UINT8', 'UNDEFINED', '__class__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__int__', '__le__', '__lt__', '__members__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__setstate__', '__sizeof__', '__str__', '__subclasshook__', 'name']

>>> dir(torch._C._onnx.OperatorExportTypes)

['ONNX', 'ONNX_ATEN', 'ONNX_ATEN_FALLBACK', 'RAW', '__class__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__int__', '__le__', '__lt__', '__members__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__setstate__', '__sizeof__', '__str__', '__subclasshook__', 'name']

initJITBindings則是通過pybind11往torch._C上注冊了一堆和JIT相關(guān)的C++函數(shù)/對象；

initNNFunctions初始化了一個torch._C._nn 對象，并注冊了一些nn相關(guān)的函數(shù)：

>>> dir(torch._C._nn)

['__doc__', '__loader__', '__name__', '__package__', '__spec__', '_parse_to', 'adaptive_avg_pool2d', 'adaptive_avg_pool3d', 'adaptive_max_pool2d', 'adaptive_max_pool3d', 'avg_pool2d', 'avg_pool3d', 'binary_cross_entropy', 'elu', 'elu_', \

'fractional_max_pool2d', 'glu', 'hardtanh', 'hardtanh_', 'l1_loss', 'leaky_relu', 'leaky_relu_', 'log_sigmoid', 'max_pool2d_with_indices', 'max_pool3d_with_indices', 'max_unpool2d', 'max_unpool3d', 'mse_loss', 'multi_margin_loss', \

'multilabel_margin_loss', 'nll_loss', 'nll_loss2d', 'reflection_pad1d', 'reflection_pad2d', 'replication_pad1d', 'replication_pad2d', 'replication_pad3d', 'rrelu_with_noise', 'rrelu_with_noise_', 'smooth_l1_loss', 'soft_margin_loss', \

'softplus', 'softshrink', 'thnn_conv2d', 'thnn_conv3d', 'thnn_conv_depthwise2d', 'thnn_conv_dilated2d', 'thnn_conv_dilated3d', 'thnn_conv_transpose2d', 'thnn_conv_transpose3d', 'upsample_bilinear2d', 'upsample_linear1d', 'upsample_nearest1d', \

'upsample_nearest2d', 'upsample_nearest3d', 'upsample_trilinear3d']

init_legacy_variable注冊了torch._C._LegacyVariableBase：

>>> dir(torch._C._LegacyVariableBase)

['__class__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', \

'__ge__', '__getattribute__', '__gt__', '__hash__', '__init__', '__le__', \

'__lt__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', \

'__setattr__', '__sizeof__', '__str__', '__subclasshook__']

_LegacyVariableBase類會派生出Variable類(該類的_execution_engine會初始化為torch._C._EngineBase)：

class Variable(with_metaclass(VariableMeta, torch._C._LegacyVariableBase))

init_bindings是通過pybind11往torch._C上注冊一些函數(shù)，torch::cuda::initModule類似，也是通過pybind11往torch._C上注冊一些函數(shù)，只不過內(nèi)容是和cuda相關(guān)的。

5，在torch._C上注冊StorageBase類

PyObject* initModule() {

......

THPDoubleStorage_init(module);

THPFloatStorage_init(module);

THPHalfStorage_init(module);

THPLongStorage_init(module);

THPIntStorage_init(module);

THPShortStorage_init(module);

THPCharStorage_init(module);

THPByteStorage_init(module);

THCPDoubleStorage_init(module);

THCPFloatStorage_init(module);

THCPHalfStorage_init(module);

THCPLongStorage_init(module);

THCPIntStorage_init(module);

THCPShortStorage_init(module);

THCPCharStorage_init(module);

THCPByteStorage_init(module);

THCPStream_init(module);

......

}

這些初始化工作主要就是往torch._C上注冊了以下類：

CudaByteStorageBase

CudaCharStorageBase

CudaDoubleStorageBase

CudaFloatStorageBase

CudaHalfStorageBase

CudaIntStorageBase

CudaLongStorageBase

CudaShortStorageBase

ByteStorageBase

CharStorageBase

DoubleStorageBase

FloatStorageBase

HalfStorageBase

IntStorageBase

LongStorageBase

ShortStorageBase

比如以FloatStorageBase為例的話，我們可以這樣查看它注冊的方法：

>>> dir(torch._C.FloatStorageBase)

['__class__', '__delattr__', '__delitem__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getitem__', '__gt__', '__hash__', '__init__', '__le__', '__len__', '__lt__', \

'__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__setitem__', '__sizeof__', '__str__', '__subclasshook__', '_cdata', '_expired', '_free_weak_ref', \

'_get_shared_fd', '_new_shared_fd', '_new_shared_filename', '_new_using_fd', '_new_using_filename', '_new_with_file', '_new_with_weak_ptr', '_set_cdata', '_set_from_file', '_share_fd_', \

'_share_filename_', '_shared_decref', '_shared_incref', '_weak_ref', '_write_file', 'copy_', 'data_ptr', 'element_size', 'fill_', 'from_buffer', 'from_file', 'is_pinned', 'is_shared', 'new', \

'resize_', 'size']

這些類會在python體系中被繼承：

class FloatStorage(_C.FloatStorageBase, _StorageBase)

另外注意下這塊代碼使用了一些宏來復(fù)用不同storage的代碼，如下所示：

aten/src/TH/THGenerateLongType.h:10:#define Real Long

aten/src/TH/THGenerateHalfType.h:10:#define Real Half

aten/src/TH/THGenerateIntType.h:10:#define Real Int

aten/src/TH/THGenerateFloatType.h:9:#define Real Float

aten/src/TH/THGenerateShortType.h:10:#define Real Short

aten/src/TH/THGenerateCharType.h:8:#define Real Char

aten/src/TH/THGenerateByteType.h:8:#define Real Byte

aten/src/TH/THGenerateDoubleType.h:9:#define Real Double

aten/src/THC/THCGenerateIntType.h:7:#define Real Int

aten/src/THC/THCGenerateLongType.h:7:#define Real Long

aten/src/THC/THCGenerateCharType.h:7:#define Real Char

aten/src/THC/THCGenerateFloatType.h:9:#define Real Float

aten/src/THC/THCGenerateDoubleType.h:7:#define Real Double

aten/src/THC/THCGenerateHalfType.h:9:#define Real Half

aten/src/THC/THCGenerateShortType.h:7:#define Real Short

aten/src/THC/THCGenerateByteType.h:7:#define Real Byte

6，ATen的初始化

本小節(jié)會進(jìn)行ATen的global context的初始化，然后使用at::globalContext().defaultGenerator(at::kCPU)進(jìn)行g(shù)enerator的初始化。

另外，PyTorch會根據(jù)編譯環(huán)境和用戶配置，然后向torch._C上注冊一些flag。這些flag有has_cudnn、has_mkl、has_lapack、_GLIBCXX_USE_CXX11_ABI：

PyObject* initModule() {

......

PyObject *has_cudnn = Py_True;

set_module_attr("has_cudnn", has_cudnn);

at::init();

py::reinterpret_borrow<:module>(module).def("_demangle", &c10::demangle);

::c10::Warning::set_warning_handler(&warning_handler);

set_module_attr("has_mkl", at::hasMKL() ? Py_True : Py_False);

set_module_attr("has_lapack", at::hasLAPACK() ? Py_True : Py_False);

set_module_attr("_GLIBCXX_USE_CXX11_ABI", _GLIBCXX_USE_CXX11_ABI ? Py_True : Py_False);

auto& defaultGenerator = at::globalContext().defaultGenerator(at::kCPU);

THPDefaultGenerator = (THPGenerator*)THPGenerator_NewWithGenerator(defaultGenerator);

set_module_attr("default_generator", (PyObject*)THPDefaultGenerator, /* incref= */ false);

7，torch._C._THNN和torch._C._THCUNN的初始化

PyTorch在這一小節(jié)里注冊了torch._C._THNN和torch._C._THCUNN類：

PyObject* initModule() {

......

torch::nn::init__THNN(module);

torch::nn::init__THCUNN(module);

......

}

這兩個類都擁有數(shù)量巨大的op函數(shù)，一個是CPU版的，一個是CUDA版的。

initModule之后

在initModule()函數(shù)初始化完畢之后，import torch的初始化工作還沒有結(jié)束。因?yàn)樵谶@之后，python的初始化腳本還要調(diào)用以下2個API才算真正完成全部的初始化：

_C._initExtension(manager_path())

_C._init_names(list(torch._storage_classes))

其中主要的工作都是在_C._initExtension中，這個初始化做了以下的工作：

torch::utils::initializeLayouts();

torch::utils::initializeDtypes();

torch::tensors::initialize_python_bindings();

THPDoubleStorage_postInit(module);

THPFloatStorage_postInit(module);

THPHalfStorage_postInit(module);

THPLongStorage_postInit(module);

THPIntStorage_postInit(module);

THPShortStorage_postInit(module);

THPCharStorage_postInit(module);

THPByteStorage_postInit(module);

THPBoolStorage_postInit(module);

//定義在THPStorage_(postInit)函數(shù)中，因?yàn)門HPStorage_會被宏替換THPDoubleStorage_ \

//THPFloatStorage_、THPHalfStorage_、THPLongStorage_......

THPAutograd_initFunctions();

最后的THPAutograd_initFunctions()則是初始化了torch的自動微分系統(tǒng)，這是PyTorch動態(tài)圖框架的基礎(chǔ)。

總結(jié)

在PyTorch的初始化階段，(python)torch模塊先后初始化產(chǎn)生torch._C、torch._C._TensorBase、pybind11綁定、torch._C.*StorageBase、torch._C._THNN、torch._C._THCUNN，并進(jìn)行了ATen context的初始化。在initModule()結(jié)束之后，初始化工作還進(jìn)行了_C._initExtension()的初始化。

總結(jié)

以上是生活随笔為你收集整理的pytorch默认初始化_PyTorch的初始化的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。