关于深度学习:OneFlow源码阅读4tensor类型体系与local-tensor

38次阅读

共计 4206 个字符,预计需要花费 11 分钟才能阅读完成。

tensor 和 op 是神经网络模型最根本的组件:op 是模型的节点,tensor 是连贯节点的边。

然而,构建一个 tensor 并不仅仅是结构一个对象那么简略,至多要思考上面这些问题:

  • 要反对节点本地的 local tensor,以及分布式的 global tensor。
  • 要反对 eager 和 lazy 执行模式。
  • 要反对不同的数据类型,包含 float、double、int 等。
  • 要反对不同设施。

1 创立 tensor 的办法

从 init.py 看,有两个办法能够创立 tensor 对象,一个是 Tensor,另一个是 tensor。这两种形式最终都会通过 PyFunction 转发到特定的 Functor。

1.1 Tensor 类型

Tensor是在 tensor.py 中引入的,构造函数被绑定为 C ++ 的 ApiNewTensor,通过 RegisterMethods 为 Tensor 注册了一些 Python 实现的办法(如将 get_item/set_item 等转发给对应的 C ++ 函数),在包初始化时会通过 RegisterMethod4Class 实现这些办法的注册。

RegisterMethod4Class 的调用流程如下:

1.2 tensor 函数

tensor 是一个函数,其绑定定义在 tensor_api.yaml.pybind.cpp 中,这是构建阶段主动生成的文件。tensor函数间接绑定到PyFunction

1.3 手动构建 tensor 的两种形式

剖析 Tensor 和 tensor 的 PyFunction 签名,能够通过如下形式结构 local tensor,也就是只能在节点外部应用的 tensor。其中只有 tensor 能够指定 dtype 参数。

import oneflow as flow
flow.tensor([[1,2,3],[4,5,6]])
flow.tensor([1, 2, 3], dtype=flow.int64)
flow.Tensor([[1,2,3],[4,5,6]])
# error
# flow.Tensor([1, 2, 3], dtype=flow.int64)

2 oneflow 的 tensor 类型体系

ApiNewTensor 函数返回 Tensor 类型。这是一个抽象类接口。通过其继承和子类的字段蕴含关系,能够失去如下的类图:

以上次要是 Tensor 相干的接口定义。MirroredTensor即节点内的 local tensor,ConsistentTensor即一致性视角的、分布式的 global tensor。

Tensor 应用了 Bridge 模式,每个 Tensor 子类外部有一个 TensorImpl 字段。TensorImpl 相干的类图如下:

3 local tensor 的结构

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

在这个例子中,TensorWithDataCtorFunctor 最终会调用 MakeLocalTensorFromData,次要的逻辑都在这个函数中。其中大量调用 Python 和 numpy 的接口,查看 PyObject 的数据类型,获取 Shape 和 DataType,如果用户没有制订 device,默认会设置为 CPU 设施。

前面次要是调用 EmptyFunctor 和 SwitchCopyMirroredTensorFromUntypedArray。前者为 tensor 分配内存,后者进行数据拷贝,两个步骤都会通过虚拟机指令实现。

为什么要通过虚拟机指令实现呢?无论是内存资源的调配,还是数据拷贝,CPU 和 CUDA 等不同设施上的操作都不一样。之前探讨 Op/Kernel 时曾经看到,虚拟机和 InstructionType 反对不同的设施,所以内存调配和数据拷贝也通过虚拟机执行。

3.1 分配内存:EmptyFunctor

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

而这里只须要结构一个 tensor 对象,不须要其它计算,所以是一个 Empty 操作,EmptyKernel 没有实质性的计算逻辑。

因为是 eager 模式下的 local tensor,EmptyFunctor 会进入 NaiveInterpret 执行。在这里会先结构 EagerMirroredTensorImpl 和 MirroredTensor 对象,用于寄存 tensor 后果。但这只是一个壳子,还没有为 tensor 的数据调配存储空间。

之后会初始化 EagerBlobObject、创立 TensorStorage,这样 tensor 次要的字段根本构建结束。

而后结构指令、提交虚拟机执行。EmptyFunctor 是 UserOp,最终会进入 LocalCallOpKernelUtil: Compute,其中 AllocateOutputBlobsMemory 实现内存分配任务。

EmptyFunctor 的调用流程如下:

AllocateOutputBlobsMemory 的调用流程如下。BlobDesc::ByteSizeOfBlobBody 提供内存 size,即 elem_cnt * SizeOf(data_type。CPU 环境下,CpuAllocator 通过aligned_alloc 申请内存资源。

3.2 拷贝数据:SwitchCopyMirroredTensorFromUntypedArray

SwitchCopyMirroredTensorFromUntypedArray其实是 MAKE_SWITCH_ENTRY 宏开展后的函数名。宏开展后的代码如下。理论会调用 CopyMirroredTensorFromUntypedArray。

template<typename... Args>
static Maybe<void> SwitchCopyMirroredTensorFromUntypedArray(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 CopyMirroredTensorFromUntypedArray<float>(std::forward<Args>(args)...);
           }},
           // ...
      };
  return case_handlers.at(switch_tuple)(std::forward<Args>(args)...);
};

数据拷贝的调用流程如下:

根据上述宏开展后的代码,CopyMirroredTensorFromUntypedArray 的模版参数是 tensor 的 dtype,如DataType::kFloat。在 tensor 结构的场景下,函数 CopyBetweenMirroredTensorAndNumpy 的模版参数如 BlobNumpyCopyUtil<DataType::kFloat>::From。

CopyBetweenMirroredTensorAndNumpy中会结构指令提交虚拟机执行。PhysicalRun的逻辑相似如下代码:

    vm::InstructionMsgList instruction_list;
    InstructionsBuilder instructions_builder(std::make_shared<vm::PhysicalIdGenerator>(),
                                            &instruction_list);
    // JUST(Build(&instructions_builder));
    builder->AccessBlobByCallback(
        tensor,
        [array_ptr, Copy](uint64_t ofblob_ptr) {CHECK_JUST(Copy(ofblob_ptr, array_ptr)); },
        modifier);
    JUST(vm::Run(instructions_builder.mut_instruction_list()));

lambda 表达式中的 Copy 就是 BlobNumpyCopyUtil<DataType::kFloat>::From;array_ptr示意 Python 端传过来的数组数据指针;前面咱们会看到,ofblob_ptr就是 tensor 的 Blob 中的指针。

InstructionsBuilder::AccessBlobByCallback中创立 AccessBlobArgCbPhyInstrOperand 对象,对应的指令类型是 AccessBlobByCallbackInstructionType。所以虚拟机执行指令时,会进入 AccessBlobByCallbackInstructionType::Compute 执行。理论的执行逻辑相似如下代码:

const auto* ptr =
  dynamic_cast<const vm::AccessBlobArgCbPhyInstrOperand*>(phy_instr_operand.get());
OfBlob ofblob(device_ctx->stream(), ptr->eager_blob_object()->mut_blob());
// ptr->callback()(reinterpret_cast<uint64_t>(&ofblob));
BlobNumpyCopyUtil<DataType::kFloat>::From(&ofblob, array_ptr);

ptr->callback()就是上述 lambda 表达式。OfBlob 是对 tensor 的 Blob 的封装。一路追踪上来,CPU 环境下最终会调用 std::memcpy 拷贝数据。

参考资料

  • oneflow v0.7.0
  • OneFlow 源码浏览 1:算子签名的主动推断
  • OneFlow 源码浏览 2:Op、Kernel 与解释器
  • OneFlow 源码浏览 3:Op 指令在虚拟机中的执行
  • 一个 Tensor 在深度学习框架中的执行过程简略梳理

正文完
 0