共计 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
matmul
和 relu
(inplace=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 在深度学习框架中的执行过程简略梳理