这篇是本系列最后一篇博客了,介绍一下前面的C++代码怎么与Python交互,或者说Python里怎么调用C++代码进行高效的计算。首先简单介绍一下预备知识,既Python的C扩展通常怎么写;然后以比较核心的数据结构 Tensor 和 Storage 为例看一下它们怎么转换为Python类型的;最后稍带点儿Python自微分函数的实现。
目录
Python的C/C++扩展 扩展模块 对于简单的C代码,构建一个自定义扩展模块是很容易的。C/C++
部分基本上只需要做以下几件事:
举个🌰,下面的代码构建了一个只有一个函数的Python模块,该函数的功能是求最大公约数(py_gcd
):
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 #include "Python.h" #include "sample.h" static PyObject *py_gcd (PyObject *self, PyObject *args) { int x, y, result; if (!PyArg_ParseTuple(args,"ii" , &x, &y)) { return NULL ; } result = gcd(x,y); return Py_BuildValue("i" , result); } static PyMethodDef SampleMethods[] = { {"gcd" , py_gcd, METH_VARARGS, "Greatest common divisor" }, { NULL , NULL , 0 , NULL } }; static struct PyModuleDef samplemodule = { PyModuleDef_HEAD_INIT, "sample" , "A sample module" , -1 , SampleMethods }; PyMODINIT_FUNC PyInit_sample (void ) { return PyModule_Create(&samplemodule); }
具体细节就不展开了,有兴趣的童鞋可以参考下面的链接。
要绑定这个扩展模块,像下面这样创建一个setup.py
文件:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 from distutils.core import setup, Extensionsetup(name='sample' , ext_modules=[ Extension('sample' , ['pysample.c' ], include_dirs = ['/some/dir' ], define_macros = [('FOO' ,'1' )], undef_macros = ['BAR' ], library_dirs = ['/usr/local/lib' ], libraries = ['sample' ] ) ] )
为了构建最终的函数库,只需简单的使用python3 buildlib.py build_ext --inplace
命令即可。它会创建一个名字叫sample.so
的共享库,当被编译后,你就能将它作为一个模块导入进来了:
1 2 3 4 >>> import sample>>> sample.gcd(35 , 42 )7 >>>
注:这部分主要参考Python3-Cookbook .
自定义Python类型 在Python代码中如果要创建一个自定义类使用class
关键字即可,但是在C代码中就没那么方便了。首先简单介绍下Python中的类型。在Python中一切皆对象,Python中有两种对象:
一种是类型对象(class对象):表示Python定义的类型,例如int, str, object
等;
另一种是实例对象(instance对象):表示由class对象创建的实例。
Python中的所有对象都是直接或者间接继承object
,然后object
又是type
类型。Python对象的C语言实现也是分为两部分,一部分表示实例对象,存储对象实际的数据;另一部分是类型对象,存储对象的元数据。也就是说,自定义类型也要实现这两部分,举个例子:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 typedef struct { PyObject_HEAD int value; } noddy_NoddyObject; static PyTypeObject noddy_NoddyType = { PyVarObject_HEAD_INIT(NULL , 0 ) "noddy.Noddy" , sizeof (noddy_NoddyObject), 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , Py_TPFLAGS_DEFAULT, "Noddy objects" , };
然后创建一个新扩展模块,并完成初始化:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 static PyMethodDef noddy_methods[] = { {NULL } }; PyMODINIT_FUNC initnoddy (void ) { PyObject* m; noddy_NoddyType.tp_new = PyType_GenericNew; if (PyType_Ready(&noddy_NoddyType) < 0 ) return ; m = Py_InitModule3("noddy" , noddy_methods, "Example module" ); Py_INCREF(&noddy_NoddyType); PyModule_AddObject(m, "Noddy" , (PyObject*)&noddy_NoddyType); }
注:这部分介绍的比较简略,详细请参考 http://www.xefan.com/archives/84091.html .
torch._C 有了上面的预备知识之后,我们就能看ATen还有autograd等模块的代码是怎么导入Python了。构建Python扩展模块的代码在torch/csrc/Module.cpp
里,主要部分如下:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 ... static std ::vector <PyMethodDef> methods;PyObject* initModule () { HANDLE_TH_ERRORS THInferNumThreads(); THPUtils_addPyMethodDefs(methods, TorchMethods); THPUtils_addPyMethodDefs(methods, DataLoaderMethods); THPUtils_addPyMethodDefs(methods, torch::autograd::python_functions()); THPUtils_addPyMethodDefs(methods, torch::multiprocessing::python_functions()); ... #if PY_MAJOR_VERSION == 2 ASSERT_TRUE(module = Py_InitModule("torch._C" , methods.data())); #else static struct PyModuleDef torchmodule = { PyModuleDef_HEAD_INIT, "torch._C" , nullptr , -1 , methods.data() }; ASSERT_TRUE(module = PyModule_Create(&torchmodule)); #endif ASSERT_TRUE(THPWrapper_init(module )); ... torch::autograd::initNNFunctions(module ); torch::autograd::init_legacy_variable(module ); torch::python::init_bindings(module ); #ifdef USE_CUDA torch::cuda::initModule(module ); #endif ASSERT_TRUE(THPDoubleStorage_init(module )); ASSERT_TRUE(THPFloatStorage_init(module )); ASSERT_TRUE(THPHalfStorage_init(module )); ASSERT_TRUE(THPLongStorage_init(module )); ASSERT_TRUE(THPIntStorage_init(module )); ASSERT_TRUE(THPShortStorage_init(module )); ASSERT_TRUE(THPCharStorage_init(module )); ASSERT_TRUE(THPByteStorage_init(module )); ASSERT_TRUE(THPBoolStorage_init(module )); ... return module ; END_HANDLE_TH_ERRORS }
基本上所有C/C++实现的API都被绑定在torch._C
扩展模块中,下面以Storage和Tensor为例,看一下torch.Storage
和torch.Tensor
类型的绑定方法,比较有意思的是它们两个的绑定方式区别还挺大的。
Storage 从上面的initModule()
函数中可以看到,里面有许多初始化各种Storage类型的代码,它们的目的就是创建各种Storage类型,如torch._C._FloatStorageBase
,torch._C._LongStorageBase
等,而torch.FloatStorage
等类型是从Python端创建的,继承自torch._C._FloatStorageBase
等类型,这部分代码可以在torch/__init__.py
中找到。
回到绑定过程,THPDoubleStorage_init()
等函数其实是用C范式生成的,和第一篇 里的TH库中用的方法一样。它实际调用的函数是bool THPStorage_(init)(PyObject *module)
,实现torch/csrc/generic/Storage.cpp
里,这个函数会根据不同类型展开:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 bool THPStorage_ (init) (PyObject *module ) { static std ::vector <PyMethodDef> methods; THPUtils_addPyMethodDefs(methods, THPStorage_(methods)); #ifndef THD_GENERIC_FILE THPUtils_addPyMethodDefs(methods, THPStorage_(sharingMethods)); #endif THPStorageType.tp_methods = methods.data(); THPStorageType.tp_members = THPStorage_(members); if (PyType_Ready(&THPStorageType) < 0 ) return false ; Py_INCREF(&THPStorageType); PyModule_AddObject(module , THPStorageBaseStr, (PyObject*)&THPStorageType); THPStorage_(initCopyMethods)(); return true ; }
这个函数初始化了THPStorageType
类型并添加到torch._C
模块中,该类型的定义如下:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 struct THPStorage { PyObject_HEAD THWStorage *cdata; }; PyTypeObject THPStorageType = { PyVarObject_HEAD_INIT(nullptr , 0 ) "torch._C." THPStorageBaseStr, sizeof (THPStorage), 0 , (destructor)THPStorage_(dealloc), ... THPStorage_(pynew), };
Python里类型的名字由tp_name
域确定,也就是"torch._C." THPStorageBaseStr
,后者一看就是宏定义,它展开后会变成 _xxxStorageBase
,其中xxx
为各种类型,所以最后就变成了 torch._C._FloatStorageBase
等类型。
Tensor torch.Tensor
的实现就与torch.Storage
不一样了,因为ATen的存在,绑定的话也是绑定ATen里的Tensor,不会绑定THTensor。还有一点,就是Python里的Tensor和Variable合并了,所以torch.Tensor
直接和 autograd::Variable
绑定在一起了。不过准确来说是 torch._C._TensorBase
和 autograd::Variable
绑定在一起了,而 torch.Tensor
继承自 torch._C._TensorBase
。
绑定的代码在 torch/csrc/autograd/python_variable.cpp
中:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 struct THPVariable { PyObject_HEAD torch::autograd::Variable cdata; PyObject* backward_hooks = nullptr ; }; ... PyTypeObject THPVariableType = { PyVarObject_HEAD_INIT(nullptr , 0 ) "torch._C._TensorBase" , sizeof (THPVariable), 0 , (destructor)THPVariable_dealloc, ... THPVariable_pynew }; bool THPVariable_initModule (PyObject *module ) { static std ::vector <PyMethodDef> methods; THPUtils_addPyMethodDefs(methods, torch::autograd::variable_methods); THPUtils_addPyMethodDefs(methods, extra_methods); THPVariableType.tp_methods = methods.data(); if (PyType_Ready(&THPVariableType) < 0 ) return false ; Py_INCREF(&THPVariableType); PyModule_AddObject(module , "_TensorBase" , (PyObject *)&THPVariableType); torch::autograd::initTorchFunctions(module ); return true ; }
注意到,这里只绑定了 _TensorBase
一种类型,而不像Storage那样利用宏把各种类型的StorageBase都定义了。
其他类型的Tensor,如 torch.FloatTensor
等在 torch/tensor/python_tensor.cpp
中的 initialize_python_bindings()
函数里动态绑定:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 void initialize_python_bindings () { initialize_aten_types(tensor_types); py_initialize_metaclass(metaclass); auto tensor_dict = get_tensor_dict(); for (auto & tensor_type : tensor_types) { py_initialize_tensor_type(tensor_type.py_type, tensor_type.name, tensor_dict.get()); } py_bind_tensor_types(tensor_types); set_default_tensor_type(at::globalContext().getVariableType( at::Backend::CPU, at::kFloat)); }
这个函数由Python调用,调用的代码在 torch/__init__.py
中(_C._initExtension()
)。调用时 torch.Tensor
已经定义,这个函数要做的就是定义其他Tensor类型,然后把Tensor类型的方法直接拷贝给它们,最后在设置一下默认类型的Tensor。为什么数据类型不同却可以直接拷贝?因为 at::Tensor
可以针对不同数据类型调用不同的方法,类型多态已经在ATen内部实现了。
在上面的代码中,Tensor 绑定的方法来自 torch::autograd::variable_methods
,这个列表在 csrc/autograd/generated/python_variable_methods.cpp
中:
1 2 3 4 5 6 7 8 9 10 11 PyMethodDef variable_methods[] = { {"__add__" , (PyCFunction)THPVariable_add, METH_VARARGS | METH_KEYWORDS, NULL }, {"__radd__" , (PyCFunction)THPVariable_add, METH_VARARGS | METH_KEYWORDS, NULL }, {"__iadd__" , (PyCFunction)THPVariable_add_, METH_VARARGS | METH_KEYWORDS, NULL }, {"__rmul__" , (PyCFunction)THPVariable_mul, METH_VARARGS | METH_KEYWORDS, NULL }, {"__mul__" , (PyCFunction)THPVariable_mul, METH_VARARGS | METH_KEYWORDS, NULL }, {"__imul__" , (PyCFunction)THPVariable_mul_, METH_VARARGS | METH_KEYWORDS, NULL }, {"__sub__" , (PyCFunction)THPVariable_sub, METH_VARARGS | METH_KEYWORDS, NULL }, {"addcmul" , (PyCFunction)THPVariable_addcmul, METH_VARARGS | METH_KEYWORDS, NULL } ... }
从文件路径可以看出这也是根据 derivatives.yaml
自动生成的代码,拿 addcmul
举个例子看看这些函数的实现方法:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 static PyObject * THPVariable_addcmul (PyObject* self_, PyObject* args, PyObject* kwargs) { HANDLE_TH_ERRORS static PythonArgParser parser ({ "addcmul(Scalar value, Tensor tensor1, Tensor tensor2)|deprecated" , "addcmul(Tensor tensor1, Tensor tensor2, *, Scalar value=1)" , }, true ) ; auto & self = reinterpret_cast <THPVariable*>(self_)->cdata; ParsedArgs<4 > parsed_args; auto r = parser.parse(args, kwargs, parsed_args); if (r.idx == 0 ) { return wrap(dispatch_addcmul(self, r.scalar(0 ), r.tensor(1 ), r.tensor(2 ))); } else if (r.idx == 1 ) { return wrap(dispatch_addcmul(self, r.tensor(0 ), r.tensor(1 ), r.scalar(2 ))); } Py_RETURN_NONE; END_HANDLE_TH_ERRORS }
最终函数调用了 dispatch_addcmul()
进行下一步计算,该函数在同文件夹的 python_variable_methods_dispatch.h
,可见这个文件也是自动生成的,函数声明如下:
1 2 3 4 5 6 inline Tensor dispatch_addcmul (Tensor & self, Scalar value, const Tensor & tensor1, const Tensor & tensor2) { AutoNoGIL no_gil; return self.addcmul(tensor1, tensor2, value); }
这个函数首先获取了GIL锁,然后调用C++前端 Variable.addcmul()
进行计算,由于后者实现了自动微分,所以Python调用也具有自动微分功能。为什么要释放GIL锁?因为这样才能使其与Python解释器中的其他进程一起正确的执行。
也就是说,Python Tensor的方法的封装思路是:
生成dispatch系列函数,该函数用于释放GIL锁,然后调用Variable的对应实现
生成可被Python调用的API,该函数解析Python参数,并调用dispatch系列函数进行实际计算
NN 神经网络的部分函数也有部分函数是直接从 ATen 绑定而来,绑定到 torch._C._nn
模块中,代码在 csrc/autograd/generate/python_nn_functions.cpp
中:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 static PyMethodDef nn_functions[] = { {"_parse_to" , (PyCFunction)THPVariable__parse_to, METH_VARARGS | METH_KEYWORDS, nullptr }, {"adaptive_avg_pool2d" , (PyCFunction)THPVariable_adaptive_avg_pool2d, METH_VARARGS | METH_KEYWORDS, NULL }, {"adaptive_avg_pool3d" , (PyCFunction)THPVariable_adaptive_avg_pool3d, METH_VARARGS | METH_KEYWORDS, NULL }, {"adaptive_max_pool2d" , (PyCFunction)THPVariable_adaptive_max_pool2d, METH_VARARGS | METH_KEYWORDS, NULL }, {"adaptive_max_pool3d" , (PyCFunction)THPVariable_adaptive_max_pool3d, METH_VARARGS | METH_KEYWORDS, NULL }, {"avg_pool2d" , (PyCFunction)THPVariable_avg_pool2d, METH_VARARGS | METH_KEYWORDS, NULL }, {"avg_pool3d" , (PyCFunction)THPVariable_avg_pool3d, METH_VARARGS | METH_KEYWORDS, NULL }, {"binary_cross_entropy" , (PyCFunction)THPVariable_binary_cross_entropy, METH_VARARGS | METH_KEYWORDS, NULL }, {"elu" , (PyCFunction)THPVariable_elu, METH_VARARGS | METH_KEYWORDS, NULL }, {"elu_" , (PyCFunction)THPVariable_elu_, METH_VARARGS | METH_KEYWORDS, NULL }, ... } void initNNFunctions(PyObject* module ) {#if PY_MAJOR_VERSION == 2 PyObject* nn = Py_InitModule("torch._C._nn" , nn_functions); Py_XINCREF(nn); #else static struct PyModuleDef def = { PyModuleDef_HEAD_INIT, "torch._C._nn" , NULL , -1 , nn_functions }; PyObject* nn = PyModule_Create(&def); #endif if (!nn) { throw python_error(); } if (PyModule_AddObject(module , "_nn" , nn) != 0 ) { throw python_error(); } }
这里面有pooling, loss, conv等相关函数,供Python里的nn.functional
调用。
此外,为了方便在Python中自定义的自微分的函数,Python里也实现了上一篇对应的Function:torch.autograd.Function
。继承它,重载 forward()
和 backward()
方法就可以用Python实现自定义的自微分函数,详见文档 。
torch.autograd.Function
的定义在 torch/autograd/function.py
中:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 class Function (with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin) ) : __call__ = _C._FunctionBase._do_forward is_traceable = False @staticmethod def forward (ctx, *args, **kwargs) : raise NotImplementedError @staticmethod def backward (ctx, *grad_outputs) : raise NotImplementedError
它继承自 torch._C.FunctionBase
,该类型在 csrc/autograd/python_function.cpp
中绑定,实现的逻辑和 autograd::Function
一样,也是在前向计算时建立反向计算图,这里就不读赘述了。绑定部分的代码如下:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 struct THPFunction { PyObject_HEAD PyObject *needs_input_grad; PyObject *to_save; PyObject *non_differentiable; PyObject *dirty_tensors; std ::vector <torch::autograd::VariableInfo> output_info; std ::vector <torch::autograd::VariableInfo> input_info; std ::vector <torch::autograd::SavedVariable> saved_variables; std ::vector <bool > is_variable_input; char has_freed_buffers; torch::autograd::PyFunction cdata; }; static struct PyGetSetDef THPFunction_properties [] = { {"saved_tensors" , (getter)THPFunction_saved_tensors, nullptr , nullptr , nullptr }, {"saved_variables" , (getter)THPFunction_saved_variables, nullptr , nullptr , nullptr }, {"next_functions" , (getter)THPFunction_next_functions, nullptr , nullptr , nullptr }, ... {nullptr } }; static struct PyMethodDef THPFunction_methods [] = { {(char *)"apply" , (PyCFunction)THPFunction_apply, METH_CLASS | METH_VARARGS, nullptr }, {(char *)"_do_forward" , (PyCFunction)THPFunction_do_forward, METH_VARARGS, nullptr }, {(char *)"_do_backward" , (PyCFunction)THPFunction_do_backward, METH_VARARGS, nullptr }, {(char *)"_register_hook_dict" , (PyCFunction)THPFunction__register_hook_dict, METH_O, nullptr }, {(char *)"register_hook" , (PyCFunction)THPFunction_register_hook, METH_O, nullptr }, {nullptr } }; PyTypeObject THPFunctionType = { PyVarObject_HEAD_INIT(nullptr , 0 ) "torch._C._FunctionBase" , sizeof (THPFunction), 0 , ... }; bool THPFunction_initModule (PyObject *module ) { if (PyType_Ready(&THPFunctionType) < 0 ) return false ; Py_INCREF(&THPFunctionType); PyModule_AddObject(module , "_FunctionBase" , (PyObject *)&THPFunctionType); return true ; }
总结一下:
完