OneFlow源码解析：Tensor类型体系与Local Tensor

撰文｜郑建华

更新｜赵露阳

tensor和op是神经网络模型最基本的组件：op是模型的节点，tensor是连接节点的边。然而，构建一个tensor并不仅仅是构造一个对象那么简单，至少要考虑以下问题：

• 要支持节点本地的local tensor，以及分布式的global tensor；

• 要支持eager和lazy执行模式；

• 要支持不同的数据类型，包括float、double、int等；

• 要支持不同设备。

1 

创建tensor的方法

与PyTorch类似，在OneFlow中也可以通过两种主要的方式来创建tensor：Tensor和tensor。这两种方式最终都会创建出OneFlow内部的C++ Tensor对象，即对应Python层的flow.Tensor类型。

1.1 Tensor

Python层的Tensor是在tensor.py（https://github.com/Oneflow-Inc/oneflow/blob/2e6a72c8734b9929191306df35b4284e9caa8126/python/oneflow/framework/tensor.py#L23）中引入的，通过python c api注册的Tensor类型对象，此对象在MakeTensorType

（https://github.com/Oneflow-Inc/oneflow/blob/2e6a72c8734b9929191306df35b4284e9caa8126/oneflow/api/python/framework/tensor.cpp#L623）中被定义和返回。

在MakeTensorType中主要通过PyTensorObject_init创建了Tensor对象：

static int PyTensorObject_init(PyObject* self, PyObject* args, PyObject* kwargs) {
  HANDLE_ERRORS
  auto* temp = functional::_legacy_tensor_ctor(NULL, args, kwargs);
  if (PyErr_Occurred()) { throw py::error_already_set(); }
  auto* _self = (PyTensorObject*)self;
  _self->data = PyTensor_Unpack(temp);
  _self->data->set_pyobject(self);
 
 
 
 
  // reset temp data to prevent clearing the pyobject
  // when the temp is deallocated
  ((PyTensorObject*)temp)->data.reset();
  Py_XDECREF(temp);
  return 0;
  END_HANDLE_ERRORS_RET(-1)
}

通过functional::_legacy_tensor_ctor函数创建了OneFlow内部的c++ Tensor对象：oneflow::one::Tensor，并作为data绑定至Python的Tensor类型。在MakeTensorType中，还通过PyMethodDef（https://github.com/Oneflow-Inc/oneflow/blob/2e6a72c8734b9929191306df35b4284e9caa8126/oneflow/api/python/framework/tensor.cpp#L639-L641）为Tensor注册了很多C++方法，如：

static PyMethodDef PyTensorObject_methods[] = {
    {"storage_offset", PyTensorObject_storage_offset, METH_NOARGS, NULL},
    {"stride", PyTensorObject_stride, METH_NOARGS, NULL},
    {"is_contiguous", PyTensorObject_is_contiguous, METH_NOARGS, NULL},
    {"contiguous", PyTensorObject_contiguous, METH_NOARGS, NULL},
    {"contiguous_", PyTensorObject_contiguous_, METH_NOARGS, NULL},
    {"pin_memory", PyTensorObject_pin_memory, METH_NOARGS, NULL},
    {"is_pinned", PyTensorObject_is_pinned, METH_NOARGS, NULL},
    {"requires_grad_", (PyCFunction)PyTensorObject_requires_grad_, METH_VARARGS | METH_KEYWORDS,
     NULL},
    {"retain_grad", PyTensorObject_retain_grad, METH_NOARGS, NULL},
    {"detach", PyTensorObject_detach, METH_NOARGS, NULL},
    {"clone", PyTensorObject_clone, METH_NOARGS, NULL},
    {"zero_", PyTensorObject_zero_, METH_NOARGS, NULL},
    {"register_hook", PyTensorObject_register_hook, METH_O, NULL},
    {"_register_post_grad_accumulation_hook", PyTensorObject__register_post_grad_accumulation_hook,
     METH_O, NULL},
    {"global_id", PyTensorObject_global_id, METH_NOARGS, NULL},
    {"check_meta_consistency", PyTensorObject_check_meta_consistency, METH_NOARGS, NULL},
    {"to_numpy", PyTensorObject_to_numpy, METH_NOARGS, NULL},
    {"type", (PyCFunction)PyTensorObject_type, METH_VARARGS | METH_KEYWORDS, NULL},

此外，在Python层通过RegisterMethods（https://github.com/Oneflow-Inc/oneflow/blob/2e6a72c8734b9929191306df35b4284e9caa8126/python/oneflow/framework/te