乐趣区

关于深度学习:OneFlow源码解析Tensor类型体系与Local-Tensor

jiezi

撰文|郑建华
更新|赵露阳

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-In…）中引入的，通过 python c api 注册的 Tensor 类型对象，此对象在 MakeTensorType
（https://github.com/Oneflow-In…）中被定义和返回。

在 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-In…）为 Tensor 注册了很多 C ++ 办法，如：

ods[] = {{"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-In…）也为 Tensor 注册了一些 Python 实现的 Tensor 办法或属性（如 tensor.numpy），在 OneFlow 包初始化时会通过 RegisterMethod4Class
（https://github.com/Oneflow-In…）实现这些 Python 办法和属性的注册。RegisterMethod4Class 的调用流程如下：

相比于 Python 实现来说，Tensor 的 ++ 实现的办法 / 属性通常具备较高的性能。

1.2 tensor 函数

Tensor 是类型，而 tensor 则是函数，flow.tensor函数在 `
oneflow/api/python/functional/tensor_api.yaml
` 中被定义：

- name: "tensor"
  signature: [
      "Tensor (PyObject* data, *, DataType dtype=None, Device device=None,
      Bool requires_grad=False, Bool pin_memory=False) => TensorWithData","Tensor (PyObject* data, *, DataType dtype=None, Placement placement,
      SbpList sbp, Bool requires_grad=False) => GlobalTensorWithData",
    ]
  bind_python: True

其 C ++ 实现位于 tensor_api.yaml.pybind.cpp 中，这是构建阶段主动生成的文件。

通过函数签名能够看到，flow.tensor()有两种重载的办法：

  • TensorWithData
  • GlobalTensorWithData

它们别离用于结构 local tensor 和 global tensor 的结构。和下面的 Tensor 相似，flow.tensor 返回的也是 OneFlow 外部的 oneflow::one::Tensor 对象（绑定至 Python 的 Tensor 对象）。

1.3 手动构建 tensor 的两种形式

和 PyTorch 相似，在 OneFlow 中罕用创立 tensor 的形式也分为两种：

  • flow.Tensor
  • flow.tensor

创立形式示例：

import oneflow
import numpy as np

oneflow.tensor([[1., -1.], [1., -1.]])
# tensor([[1., -1.],
#         [1., -1.]], dtype=oneflow.float32)
oneflow.tensor(np.array([[1, 2, 3], [4, 5, 6]]))
# tensor([[1, 2, 3],
#         [4, 5, 6]], dtype=oneflow.int64)
flow.Tensor([[1,2,3],[4,5,6]])

大多数状况下（和 PyTorch 相似的 eager 模式），能够通过指定 device、dtype、shape 等参数创立一般 tensor（local tensor）；

多数状况下（如 OneFlow 特有的 eager global、lazy 模式），须要 global tensor 时，能够通过指定 sbp 和 placement 的形式间接创立 global tensor，也可通过 tensor.to_global 的形式将一般 tensor 转换为 global tensor，可参考：

  • oneflow.tensor（https://oneflow.readthedocs.i…）
  • global tensor
    （https://docs.oneflow.org/mast…）

2

OneFlow 的 tensor 类型体系

上述内容中介绍的 oneflow 外部的 C ++ Tensor 对象，实际上其定义位于：oneflow/core/framework/tensor.h，是一个形象的 Tensor 类型。

其中 LocalTensor 即为一般的单卡视角下的 Tensor（和 PyTorch 的 Tensor 相似）；GlobalTensor则为 OneFlow 所特有的全局视角下的 Tensor（通常用于 eager global 模式或 lazy 模式下）。Tensor 应用了 Bridge 模式，每个 Tensor 子类外部有一个 TensorImpl 字段，负责形象 Tensor 的理论实现：

3

local tensor 的结构

咱们以 flow.tensor([[1,2,3],[4,5,6]])为例，看一下 tensor 结构的过程。次要的流程如下：

在这个例子中，因为应用的是 flow.tensor 办法创立 tensor（且为一般的 local tensor）所以会用到在 `
oneflow/api/python/functional/tensor_api.yaml
中定义的 TensorWithData 办法，其实现，是位于
oneflow/api/python/functional/tensor_api.cpp
` 的 TensorWithDataFunctor：

class TensorWithDataFunctor {
 public:
  Maybe<Tensor> operator()(PyObject* data, const Optional<Symbol<DType>>& dtype,
                           const Optional<Symbol<Device>>& device, const bool requires_grad,
                           const bool pin_memory) const {
    ...
    if (PyTensor_Check(data)) {
      // Throw warnings like pytorch.
      auto ret = PyErr_WarnEx(
          PyExc_UserWarning,
          "To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach()"
          "or sourceTensor.clone().detach().requires_grad_(True), rather than"
          "oneflow.tensor(sourceTensor).",
          1);
      if (ret != 0) {return Error::RuntimeError(); }

      const auto& other = PyTensor_Unpack(data);
      return MakeTensorFromOtherTensor(other, dtype, device, requires_grad, pin_memory);
    } else {
      // Make tensor from python sequence or numpy array.
      return MakeLocalTensorFromData(data, dtype, device, requires_grad, pin_memory);
    }
  }
};

因为这里传入的 data 是一个 Python 的 list 对象，所以最终会调用 MakeLocalTensorFromData 办法，创立 tensor 次要的逻辑都在这个函数中。其中大量调用 Python 和 Numpy 的接口，查看 PyObject 的数据类型，获取 Shape
（https://github.com/Oneflow-In…）和 DataType（https://github.com/Oneflow-In…），如果用户没有制订 device，默认会设置为 CPU 设施（https://github.com/Oneflow-In…）。

前面次要是调用 EmptyFunctor
（https://github.com/Oneflow-In…）和 SwitchCopyLocalTensorFromUntypedArray（https://github.com/Oneflow-In…）。前者为 tensor 分配内存，后者进行数据拷贝，两个步骤都会通过虚拟机指令实现。其中 EmptyFunctor 会走一般的 OpCall 指令、而 CopyLocalTensorFromUntypedArray 会依据是否须要同步 copy 走到 `
AccessBlobByCallback/SyncAccessBlobByCallback
` 指令。

为什么要通过虚拟机指令实现呢？无论是内存资源的调配，还是数据拷贝，CPU 和 CUDA 等不同设施上的操作都不一样。之前探讨 Op/Kernel 时曾经看到，在 OneFlow 中所有动动态图工作执行、eager 模式下 op/kernel 执行、内存 / 显存的调配和开释、device、stream 等对立由虚拟机进行治理。

3.1 分配内存：EmptyFunctor

matmul和 relu（inplace=false 时）等操作在执行过程中也会创立 output tensor。之前探讨 relu 时重点关注了 op 和 kernel 的计算逻辑，而疏忽了 tensor 相干的内容。

而这里只须要先结构一个空 tensor 对象，不须要其它计算，所以是一个 Empty 操作，Empty op 对应的 kernel——EmptyKernel（https://github.com/Oneflow-In…）没有实质性的计算逻辑，只是先依据 shape、dtype、device 信息创立一个空 tensor，期待后续将理论的数据从内存中 copy 至此空 tensor，从而实现整个 tensor 的创立过程。

EmptyFunctor 同样和其余 functor 一样，最终会被 Dispacth 至对应的 interpreter 被解释执行，这里因为是 eager 模式下的 local tensor，EmptyFunctor 最终会进入 eager local interpreter，交给 NaiveInterpret（https://github.com/Oneflow-In…）办法解决。流程如下：

  1. 在结构 EagerLocalTensorImpl（https://github.com/Oneflow-In…）对象，用于寄存 tensor 后果。但这只是一个壳子，还没有为 tensor 的数据调配存储空间。
  2. 之后会初始化 EagerBlobObject（https://github.com/Oneflow-In…）、TensorStorage（https://github.com/Oneflow-In…），这样 tensor 次要的字段根本构建结束
  3. 而后结构 OpCall 指令、提交虚拟机 PhysicalRun（https://github.com/Oneflow-In…），期待 vm 的调度执行。

OpCall 对应的指令策略最终会进入 oneflow/core/vm/op_call_instruction_policy.cpp，并在Prepare 办法中通过 AllocateOutputBlobsMemory 办法对 TensorStorage 实现理论的内存调配；在 Compute 办法中启动（empty op 对应的）理论的 kernel 执行。

3.2 拷贝数据：SwitchCopyLocalTensorFromUntypedArray

SwitchCopyMirroredTensorFromUntypedArray其实是 MAKE_SWITCH_ENTRY（https://github.com/Oneflow-In…）宏开展后的函数名。宏开展后的代码如下。理论会调用 CopyLocalTensorFromUntypedArray（https://github.com/Oneflow-In…）。

template<typename... Args>
static Maybe<void> SwitchCopyLocalTensorFromUntypedArray(const std::tuple<DataType>& switch_tuple, Args&& ... args) {static const std::map<std::tuple<DataType>, std::function<Maybe<void>(Args && ...)>>
      case_handlers {{SwitchCase(DataType::kFloat),
           [](Args&&... args) {return CopyLocalTensorFromUntypedArray<float>(std::forward<Args>(args)...);
           }},
           // ...
      };
  return case_handlers.at(switch_tuple)(std::forward<Args>(args)...);
};

CopyLocalTensorFromUntypedArray 办法如下：

template<typename T>
Maybe<void> CopyLocalTensorFromUntypedArray(const std::shared_ptr<Tensor>& tensor,
                                            PyObject* array) {
  return CopyBetweenLocalTensorAndNumpy<T>(tensor, array, CopyFromNumpyArray, "mut",
                                           /*block_host_until_done=*/false);
}

其外部理论调用了 CopyBetweenLocalTensorAndNumpy 办法。

CopyBetweenLocalTensorAndNumpy

顾名思义，这个办法次要是用在 numpy 和 tensor 之间进行数据 copy 的。其中第 3 个参数：CopyFromNumpyArray理论是一个函数回调的 callback 办法，其次要通过 SyncAutoMemcpy 进行 array 和 tensor(blob)之间的内存拷贝：

void CopyFromNumpyArray(ep::Stream* stream,
                        const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object,
                        const NumPyArrayPtr& array_ptr) {SyncAutoMemcpy(stream, eager_blob_object->mut_dptr(), array_ptr.data(),
                 eager_blob_object->ByteSizeOfBlobBody(), eager_blob_object->mem_case(),
                 memory::MakeHostMemCase());
}

持续看 CopyBetweenLocalTensorAndNumpy（https://github.com/Oneflow-In…）办法，其中最要害的是：

JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
      return builder->AccessBlobByCallback(
          tensor,
          [array_ptr, Copy](ep::Stream* stream,
                            const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object) {Copy(stream, eager_blob_object, array_ptr);
          },
          modifier);
    }));

通过 InstructionsBuilder 构建了 AccessBlobByCallback 指令，参数为下面通过 EmptyFuncor 创立的空 tensor、callback 的函数指针及参数、以及 modifier（string “mut” 示意可动静批改）。

AccessBlobByCallback

和 OpCall 相似，InstructionsBuilder 调用 AccessBlobByCallback 时，也会理论结构对应的 vm 指令策略——AccessBlobArgCbInstructionPolicy并派发至 vm，期待被调度和理论执行：

template<typename T>
Maybe<void> InstructionsBuilder::AccessBlobByCallback(
    const T tensor,
    const std::function<void(ep::Stream*, const std::shared_ptr<vm::EagerBlobObject>&)>& callback,
    const std::string& modifier) {const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object = JUST(tensor->eager_blob_object());
  Symbol<Device> device = JUST(GetDevice(tensor));
  ...
  Symbol<Stream> stream = JUST(GetDefaultStreamByDevice(device));
  JUST(SoftSyncStream({eager_blob_object}, stream));
  auto instruction = intrusive::make_shared<vm::Instruction>(
      // Never replace `stream` with producer_stream or last_used_stream.
      JUST(Singleton<VirtualMachine>::Get()->GetVmStream(stream)),
      std::make_shared<vm::AccessBlobArgCbInstructionPolicy>(eager_blob_object, callback,
                                                             modifier));
  instruction_list_->EmplaceBack(std::move(instruction));
  return Maybe<void>::Ok();}

等该条 AccessBlobArgCbInstructionPolicy 指令理论执行时，会在指令的 Compute（https://github.com/Oneflow-In…）办法中调用 callback 实现从 tensor 的 blob <-> numpy 的 ndarray 之间的数据 copy，至此拷贝过程完结，flow.tensor的创立全副实现。

（本文经受权后公布。原文：https://segmentfault.com/a/11…）

参考资料

  • On‍eFlow 源码：https://github.com/Oneflow-In…
  • OneFlow 源码解析：Op、Kernel 与解释器
  • OneFlow 源码解析：算子指令在虚拟机中的执行

欢送下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In…

退出移动版