python 建立loop Python 建立dcc-garch模型

这篇是本系列最后一篇博客了，介绍一下前面的C++代码怎么与Python交互，或者说Python里怎么调用C++代码进行高效的计算。首先简单介绍一下预备知识，既Python的C扩展通常怎么写；然后以比较核心的数据结构 Tensor 和 Storage 为例看一下它们怎么转换为Python类型的；最后稍带点儿Python自微分函数的实现。

Python的C/C++扩展

扩展模块

对于简单的C代码，构建一个自定义扩展模块是很容易的。C/C++部分基本上只需要做以下几件事：

举个例子，下面的代码构建了一个只有一个函数的Python模块，该函数的功能是求最大公约数（py_gcd）:

#include "Python.h"
#include "sample.h"

/* 把普通C语言实现的gcd()封装成Python可以调用的函数 */
static PyObject *py_gcd(PyObject *self, PyObject *args) {
  int x, y, result;

  /* 从 args 里解析实际参数 */
  if (!PyArg_ParseTuple(args,"ii", &x, &y)) {
    return NULL;
  }
  /* 调用普通C语言实现的gcd() */
  result = gcd(x,y);
  /* 把 int 转化为 PyObject* */
  return Py_BuildValue("i", result);
}

/* 定义模块的 method table */
static PyMethodDef SampleMethods[] = {
  {"gcd",  py_gcd, METH_VARARGS, "Greatest common divisor"},
  { NULL, NULL, 0, NULL}
};

/* 定义模块结构 */
static struct PyModuleDef samplemodule = {
  PyModuleDef_HEAD_INIT,
  "sample",           /* name of module */
  "A sample module",  /* Doc string (may be NULL) */
  -1,                 /* Size of per-interpreter state or -1 */
  SampleMethods       /* Method table */
};

/* 模块初始化函数 */
PyMODINIT_FUNC PyInit_sample(void) {
  return PyModule_Create(&samplemodule);
}

要绑定这个扩展模块，像下面这样创建一个setup.py文件：

# setup.py
from distutils.core import setup, Extension

setup(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的共享库，当被编译后，你就能将它作为一个模块导入进来了：

>>> import sample
>>> sample.gcd(35, 42)
7

自定义Python类型

在Python代码中如果要创建一个自定义类使用class关键字即可，但是在C代码中就没那么方便了。首先简单介绍下Python中的类型。在Python中一切皆对象，Python中有两种对象：

• 一种是类型对象（class对象）：表示Python定义的类型，例如int, str, object等；
• 另一种是实例对象（instance对象）：表示由class对象创建的实例。

Python中的所有对象都是直接或者间接继承object，然后object又是type类型。Python对象的C语言实现也是分为两部分，一部分表示实例对象，存储对象实际的数据；另一部分是类型对象，存储对象的元数据。也就是说，自定义类型也要实现这两部分，举个例子：

/* 实例对象 */
typedef struct {
    PyObject_HEAD
    /* 类型实际的数据在这里定义 */
    int value;
} noddy_NoddyObject;

/* 类型对象 */
static PyTypeObject noddy_NoddyType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "noddy.Noddy",             /*tp_name*/
    sizeof(noddy_NoddyObject), /*tp_basicsize*/
    0,                         /*tp_itemsize*/
    0,                         /*tp_dealloc*/
    0,                         /*tp_print*/
    0,                         /*tp_getattr*/
    0,                         /*tp_setattr*/
    0,                         /*tp_compare*/
    0,                         /*tp_repr*/
    0,                         /*tp_as_number*/
    0,                         /*tp_as_sequence*/
    0,                         /*tp_as_mapping*/
    0,                         /*tp_hash */
    0,                         /*tp_call*/
    0,                         /*tp_str*/
    0,                         /*tp_getattro*/
    0,                         /*tp_setattro*/
    0,                         /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT,        /*tp_flags*/
    "Noddy objects",           /*tp_doc*/
};

然后创建一个新扩展模块，并完成初始化：

/* 定义模块的 method table */
static PyMethodDef noddy_methods[] = {
    {NULL}
};

/* 模块初始化函数 */
PyMODINIT_FUNC initnoddy(void) 
{
    PyObject* m;
    /* tp_new 相当于Python里的 __new__ */
    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);
}

torch._C

有了上面的预备知识之后，我们就能看ATen还有autograd等模块的代码是怎么导入Python了。构建Python扩展模块的代码在torch/csrc/Module.cpp里，主要部分如下：

/* method table */
static std::vector<PyMethodDef> methods;

/* 初始化模块 */
PyObject* initModule() {
  HANDLE_TH_ERRORS
  THInferNumThreads();

  /* 向 method table 添加函数 */
  THPUtils_addPyMethodDefs(methods, TorchMethods);
  THPUtils_addPyMethodDefs(methods, DataLoaderMethods);
  THPUtils_addPyMethodDefs(methods, 
                           torch::autograd::python_functions());
  THPUtils_addPyMethodDefs(methods, 
                       torch::multiprocessing::python_functions());
	...

  /* 构建 torch._C 模块 */
#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);     // 初始化自微分相关API
  torch::autograd::init_legacy_variable(module);// 初始化Tensor类型
  torch::python::init_bindings(module);         // 初始化NN相关函数
#ifdef USE_CUDA
  torch::cuda::initModule(module);              // 初始化CUDA模块
#endif
  /* 初始化各种Storage类型 */
  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里，这个函数会根据不同类型展开：

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);
  /* 向模块添加 THPStorage 类型 */
  PyModule_AddObject(module, THPStorageBaseStr, 
                     (PyObject*)&THPStorageType);
  THPStorage_(initCopyMethods)();
  return true;
}

这个函数初始化了THPStorageType类型并添加到torch._C模块中，该类型的定义如下：

/* 实例对象 */
struct THPStorage {
  PyObject_HEAD
  /* THWStorage 为宏定义，会转换为 THxxxStorage */
  THWStorage *cdata;
};

/* 类型对象 */
PyTypeObject THPStorageType = {
  PyVarObject_HEAD_INIT(nullptr, 0)
  "torch._C." THPStorageBaseStr,         /* tp_name */
  sizeof(THPStorage),                    /* tp_basicsize */
  0,                                     /* tp_itemsize */
  (destructor)THPStorage_(dealloc),      /* tp_dealloc */
  ...
  THPStorage_(pynew),                    /* tp_new */
};

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中：

// python_variable.h
/* 实例对象 */
struct THPVariable {
    PyObject_HEAD
    // Payload
    torch::autograd::Variable cdata;
    PyObject* backward_hooks = nullptr;
};

// python_variable.cpp
...
/* 类型对象 */
PyTypeObject THPVariableType = {
  PyVarObject_HEAD_INIT(nullptr, 0)
  "torch._C._TensorBase",                /* tp_name */
  sizeof(THPVariable),                   /* tp_basicsize */
  0,                                     /* tp_itemsize */
  (destructor)THPVariable_dealloc,       /* tp_dealloc */
  ...
  THPVariable_pynew                      /* tp_new */
};

/* 初始化类型 */
bool THPVariable_initModule(PyObject *module)
{
  /* 获取 method table */
  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()函数里动态绑定：

void initialize_python_bindings() {
  /* 把ATen里的Tensor类型转化为Python里的PyTypeObject */
  initialize_aten_types(tensor_types);

  /* 初始化上面转化来的PyTypeObject */
  py_initialize_metaclass(metaclass);

  /* 获取 torch.Tensor 的所有方法 */
  auto tensor_dict = get_tensor_dict();

  /* 把torch.Tensor的方法复制给每个类型，如torch.FloatTensor等 */
  for (auto& tensor_type : tensor_types) {
    py_initialize_tensor_type(tensor_type.py_type, 
                              tensor_type.name, tensor_dict.get());
  }

  /* 向torch模块绑定这些各种类型的Tensor */
  py_bind_tensor_types(tensor_types);

  /* 设置 torch.Tensor 的默认类型为 torch.FloatTensor */
  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中：

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举个例子看看这些函数的实现方法：

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)",
  }, /*traceable=*/true);
  /* 获取autograd::Variable实例 */
  auto& self = reinterpret_cast<THPVariable*>(self_)->cdata;
  /* 解析函数参数 */
  ParsedArgs<4> parsed_args;
  auto r = parser.parse(args, kwargs, parsed_args);

  /* 调用 dispatch */
  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，可见这个文件也是自动生成的，函数声明如下：

inline Tensor dispatch_addcmul(Tensor & self, Scalar value, const Tensor & tensor1, const Tensor & tensor2) {
  /* 释放GIL锁 */
  AutoNoGIL no_gil;
  /* 调用 autograd::Variable.addcmul */
  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中：

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中：

class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin)):
    # only for backward compatibility
    __call__ = _C._FunctionBase._do_forward

    # for the tracer
    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一样，也是在前向计算时建立反向计算图，这里就不读赘述了。绑定部分的代码如下：

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;
    // For each input, true if the input is a THPVariable
    std::vector<bool> is_variable_input;
    char has_freed_buffers;

    /* PyFunction继承自Function，实际调用最终转发到Function中 */
    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",              /* tp_name */
  sizeof(THPFunction),                   /* tp_basicsize */
  0,                                     /* tp_itemsize */
  ...
};

bool THPFunction_initModule(PyObject *module)
{
  if (PyType_Ready(&THPFunctionType) < 0)
    return false;
  Py_INCREF(&THPFunctionType);
  PyModule_AddObject(module, "_FunctionBase", 
                     (PyObject *)&THPFunctionType);
  return true;
}

总结一下：

• THPFunction= _C._FunctionBase
• Python的autograd.Function继承自上面实现自动微分