Python 扩展 Op
阅读原文时间:2023年07月10日阅读:4

Python 扩展 Op

注意 :本文涉及的 Python Kernel 仅在 gcc 4.8.5 编译环境下充分测试,进一步的完善计划见 Issue 3951

背景介绍

OneFlow 将各种对于数据的处理都抽象成了算子(operator),简称 op。 op 是作用在输入 tensor 上的操作,并将操作的结果写到输出 tensor 上。OneFlow 内部已经提供了比较完备的 op 算子,可以在 ops 目录下找到。

当 OneFlow 已有的 Python 算子及其组合无法满足构建神经网络的需求,或者 Python 层次的算子无法满足性能需求时,可以开发自定义 op。OneFlow 提供了两类开发自定义 Op 的途径,一类是以 Python 为主的 Python Kernel 开发,另外一类是使用 C++ 扩展 Op一文介绍的 C++ Kernel 开发。

Python Kernel 因为主要采用 Python 进行扩展,开发流程较简单,适用于快速预研、算法验证等场景。C++ Kernel 效率高,适用于开发已经验证稳健性并追求性能的算子。

本文介绍算子开发的背景知识和基本概念,并展示如何开发 Python Kernel。

基本概念

在进行 OneFlow 算子开发前,需要了解 op_type_name、Op 以及 Kernel 这几个概念:

  • op_type_name:op_type_name 是 op 类别的全局唯一 ID, OneFlow 通过 op_type_name 查询并确认 op 的种类,进而实例化 op,用于构建计算图。op 的种类与 op 的关系,类似于类与对象的关系。
  • op:逻辑上的算子,包含构图推理时的输入输出形状等信息,不包含具体的处理数据的逻辑。
  • kernel:对于一个逻辑上的 op,在运行时,处理的逻辑会因为物理设备以及数据类型的不同。运行时的具体处理逻辑,由 kernel 完成。简单而言,op 与 kernel 是一对多的关系,可以使用 Python 完成具体运算,这样的Kernel 称为 Python Kernel,也可以使用 C++ 开发 Kernel
  • OneFlow 的内核由 C++ 实现,但是用户接口使用 Python,需要按照约定编写 Python Wrapper,使得 Python Op 接口能与 C++ 内核交互。

开发步骤

使用 Python 扩展 Op,应该准备一个以 op_type_name 命名的目录,在该目录下,按照约定放置必需的文件,以 oneflow/python/test/custom_ops/user_sigmoid 为例:

user_sigmoid

├── user_sigmoid_cpp_def.cpp

├── user_sigmoid_py_api.py

└── user_sigmoid_py_kernel.py

其中:

  • op_type_name_cpp_def.cpp(以上的 user_sigmoid_cpp_def.cpp) 文件中放置 Op 定义信息
  • op_type_name_py_api.py(以上的 user_sigmoid_py_api.py)文件中放置 Python Wrapper,通过 oneflow.user_op_builder 将实现的 Python Kernel 导出给用户使用
  • op_type_name_py_kernel.py(以上的 user_sigmoid_py_kernel.py)文件中放置 Python 实现的自定义算子的前向计算逻辑和后向计算逻辑

在下文中,将分别介绍:

  • 如何编写 op_type_name_cpp_def.cpp 文件,定义 Op 信息
  • 如何编写 op_type_name_py_api.py 文件,封装 Op 的 Python 接口
  • 如何编写 op_type_name_py_kernel.py 文件,使用 Python 实现 Op 的计算 Kernel
  • 在 OneFlow 中如何使用 Python Kernel 类型的自定义 Op

下文中,将介绍如何用 Python 实现一个自定义的 user_relu Op。

Op 的实现与注册

首先,在 user_relu_cpp_def.cpp 中定义 op 并完成注册:

#include "oneflow/core/framework/framework.h"

namespace oneflow {

namespace {

REGISTER_USER_OP("user_relu_forward")

.Attr("device_sub_tag", "py")

.Input("in")

.Output("out")

.SetTensorDescInferFn(

[](user_op::InferContext *ctx) -> Maybe {

*ctx->Shape4ArgNameAndIndex("out", 0) =

*ctx->Shape4ArgNameAndIndex("in", 0);

*ctx->Dtype4ArgNameAndIndex("out", 0) =

*ctx->Dtype4ArgNameAndIndex("in", 0);

return Maybe::Ok();

});

}  // namespace

}  // namespace oneflow

分析以上代码:

  • oneflow/core/framework/framework.h 中包含了创建一个 op 所需要的所有接口
  • .Attr("device_sub_tag", "py") 是必需的,它告知 OneFlow 在使用该 Op 时默认调用Python Kernel
  • 与自定义 op 有关的接口集中在 oneflow::user_op 中,使用名称空间 oneflow 可以简化类型名称
  • 宏 REGISTER_USER_OP 用于注册 op,其接受的参数 user_relu_forward 是 op_type_name。
  • 使用 REGISTER_USER_OP 注册后,其实会返回一个 OpRegistry 类(位于user_op_registry.h),通过调用该类方法,完成对自定义 op 的设置:
  1. Input("in") 表示其有一个名为 "in" 的输入
  2. Output("out") 表示其有一个名为 "out" 的输出
  3. SetTensorDescInferFn 用于设置形状及数据类型推导函数,描述该算子的输出的形状及类型与输入的关系。以上代码中,输出的形状、数据类型与输入的一致

op_type_name_cpp_def.cpp 文件是实现 Python Kernel 过程中唯一会使用到的 C++ 文件,用于设置 Op 的信息,在现阶段,还无法将使用 C++ 配置 Op 的步骤省略(因为设置分布式等高级信息时必需),不过可以看到,该文件并不涉及具体的运算,仅仅是用于描述 Op,即使不熟悉 C++,根据的示例,也可以很轻松地掌握。

封装 Op 的 Python 接口

为了用户可以在 Python 层使用刚刚设置并注册的 user_relu Op,需要创建一个 user_relu_py_api.py 文件,其内容如下:

import oneflow as flow

def user_relu_forward(x):

op = (

flow.user_op_builder("myrelu")

.Op("user_relu_forward")

.Input("in", [x])

.Output("out")

.Build()

)

return op.InferAndTryRun().SoleOutputBlob()

flow.user_op_builder("op_myrelu") 其实会返回一个名为 op_myrelu 的 UserOpConfBuilder 对象。

该对象包含 Op、Input 等方法,用于封装自定义 op,具体解释如下:

  • Op("user_relu_forward"):参数必须为之前在 C++ 注册时的 op_type_name,OneFlow 通过它找到已经注册的 op 类型,并实例化 op 对象。
  • Input("in", [input_blob]):对应了 C++ 中 op 注册时的 Input,第一个参数字符串必须与 C++ 注册 op 时的 Input 设置的字符串一致。第二个参数为输入的张量,是一个 list,因为一个 op 允许有多个输入。
  • Output("out"):对应了 C++ 中 op 注册时的 Output。
  • Build:以上设置完成后,调用 Build 可以得到自定义 op 的 Python wrapper

以下代码,将获取自定义 op 的输出:

return op.InferAndTryRun().SoleOutputBlob()

其中的 InferAndTryRun 完成推导,返回 UserOp,如果返回结果只有一个输出,则使用 SoleOutputBlob 即可获取该唯一输出,否则,可以使用 RemoteBlobList 获取包含多个输出的列表。

使用 Python 实现 Kernel

如本文开始所描述,Op 只是逻辑上的概念,真正的计算需要 Kernel 完成,在 OneFlow 中可以既可以使用 C++ ,也可以使用 Python 实现 Kernel,本文只介绍最易上手的 Python Kernel 的实现方法。使用 C++ 实现 Kernel 可以参考使用 C++ 开发 Kernel

为了为上文设置的 user_relu Op 提供 Python Kernel,需要创建一个 user_relu_py_kernel.py 文件,其内容如下:

import numpy as np

def forward(args):

(x,) = args

y = (x>0)*x

return y

以上的 forward 方法是必需实现的,它的实现对应了 Op 的 Python Kernel。关于它的约定有:

  • 方法名必需为 forward
  • 参数只有一个,类型为 tuple,tuple 中的元素个数和顺序,与 Op 注册时的 Input 对应。如之前为 user_relu 注册了 Input("in"),那么以上代码中 (x, ) = args 中的 x 就取到 in 的值
  • 输出与 Op 注册时的 Output 对应
  • 参数与返回值均为 numpy 对象,即不能(不会)是字符串、整型数字等其它类型

使用自定义 Op

完成以上工作后,得到了一个名为 user_relu 的目录,包含三个文件,它们的结构如下:

user_relu/

├── user_relu_cpp_def.cpp

├── user_relu_py_api.py

└── user_relu_py_kernel.py

可以在 user_relu 文件夹所在的路径,创建一个测试文件,调用刚刚实现的自定义 Op,内容如下:

import oneflow as flow

import numpy as np

import os

import oneflow.typing as tp

# 根据指定的路径与 op_type_name 创建 module 对象

module_path = os.path.dirname(os.path.abspath(__file__))

user_relu_op = flow.experimental.custom_op_module("user_relu", module_path)

# 使 Op, Python API, Python Kernel 生效

user_relu_op.py_api().cpp_def().py_kernel().build_load()

@flow.global_function()

def MyJob(x: tp.Numpy.Placeholder((5,), dtype=flow.float32)) -> tp.Numpy:

with flow.scope.placement("cpu", "0:0"):

return user_relu_op.api.user_relu_forward(x)

if __name__ == "__main__":

input = np.array([-2, -1, 0, 1, 2], dtype=np.float32)

output = MyJob(input)

print(input)

print(output)

以上代码中,先通过 flow.experimental.custom_op_module 创建 module 对象,它接收两个参数,第一个参数为 op_type_name, 第二个参数为 user_relu 文件夹所在的路径。返回的 module 对象,代表了自定义的 Op。

接着,通过 user_sigmoid_op.py_api().cpp_def().py_kernel().build_load() 可以使自定义 Op 生效,生效后的 Op 的 Python 接口,就是定义在 user_relu_py_api.py 文件中的方法名(user_relu_forward),它被放置在 moudle 对象的 api 名称空间中。需要通过以下方式调用:

user_sigmoid_op.api.user_relu_forward(x)

且因为 Python Kernel 只能运行在 CPU 设备上,需要指定计算设备为 CPU:

with flow.scope.placement("cpu", "0:0"):

为自定义 Op 提供反向计算

通过上述工作,已经完成了 user_relu 算子的正向计算过程,可以用于 type="predict" 的作业函数。如果想支持 type="train" 类型的训练作业函数,需要为自定义 Op 提供反向计算。

自定义 Op 提供反向计算的代码,需要写在 op_type_name_cpp_def.cpp 文件中,通过宏 REGISTER_USER_OP_GRAD 进行注册。

从数学角度上看,注册过程就是为自定义的 op,指定后向求梯度的计算方法。从编程角度看,就是为自定义 op 设置一个后向生成函数,在该函数中,编写代码,指定这个 op 的输入梯度的计算方法。

以下,将专门实现一个 Op,名为 user_relu_backward。将在为 user_relu 注册后向梯度时,用到这个“专门定制”的 Op。

实现 user_relu_backward Op

实现 user_relu_backward Op 的过程与实现 user_relu 的前向几乎是一样的。首先,在 user_relu_cpp_def.cpp 中设置并注册该 Op:

REGISTER_USER_OP("user_relu_backward")

.Input("y")

.Input("dy")

.Output("dx")

.Attr("device_sub_tag", "py")

.SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe {

const Shape* dy_shape = ctx->Shape4ArgNameAndIndex("dy", 0);

Shape* dx_shape = ctx->Shape4ArgNameAndIndex("dx", 0);

*dx_shape = *dy_shape;

return Maybe::Ok();

});

值得注意的是,同前向类似,以上代码中 .Attr("device_sub_tag", "py") 必不可少,它告知 OneFlow 在使用该 Op 时,默认调用 Python Kernel。

同理,不需要用户直接调用这个 user_relu_backward Op,不需要在 user_relu_py_api.py 为 user_relu_backward 封装 Python 接口。可以直接实现它的 Python Kernel。

在 user_relu_py_kernel.py 中,实现 backward 方法:

def backward(args):

(y, dy) = args

dx = (y>0)*dy

return dx

它的参数是一个 tuple,数目和顺序对应了 Op 注册时的 Input,输出对应了 Op 注册时的 Output。

为 Op 注册反向梯度

需要在 user_relu_cpp_def.cpp 中,通过宏 REGISTER_USER_OP_GRAD 为的正向 Op (user_relu_forward) 注册反向。

其代码如下:

REGISTER_USER_OP_GRAD("user_relu_forward")

.SetBackwardOpConfGenFn([](user_op::BackwardOpConfContext* ctx) {

const auto grad_op_name = ctx->FwOp().op_name() + "_grad";

const auto& grad_op_func = [&ctx](user_op::BackwardOpBuilder& builder) {

return builder.OpTypeName("user_relu_backward")

.InputBind("y", ctx->FwOp().output("y", 0))

.InputBind("dy", ctx->FwOp().output_grad("y", 0))

.Output("dx")

.Build();

};

ctx->DefineOp(grad_op_name, grad_op_func);

const auto& dx_get_func = [&ctx, &grad_op_name]() -> const std::string& {

return ctx->GetOp(grad_op_name).output("dx", 0);

};

ctx->FwOp().InputGradBind(user_op::OpArg("x", 0), dx_get_func);

});

对以上代码进行解释,通过 REGISTER_USER_OP_GRAD("user_relu_forward") 注册为前向 Op 注册后向求梯度规则,该宏接收一个参数,就是 前向的 op_type_name。

然后通过 SetBackwardOpConfGenFn 设置后向求梯度规则,同 Op 类似,在 op_type_name_cpp_def.cpp 中注册后向,其实不涉及真正的运算,而是设置后向计算与前向的对应关系,告诉 OneFlow 框架:

  • 用什么 Op 求后向梯度
  • 该 Op 的输入来自哪里,和前向 Op 什么关系

因此,以上代码中的:

const auto& grad_op_func = [&ctx](user_op::BackwardOpBuilder& builder) {

return builder.OpTypeName("user_relu_backward")

.InputBind("y", ctx->FwOp().output("y", 0))

.InputBind("dy", ctx->FwOp().output_grad("y", 0))

.Output("dx")

.Build();

};

定义了 Op 求梯度的方法:使用 user_relu_backward 算子,并且将前向的输出 y 作为 user_relu_backward 的输入 y;将前向的输出 y 的梯度,作为 user_relu_backward 的输入 dy;最后输出 dx。

定完求梯度的方法后,需要调用

ctx->DefineOp(grad_op_name, grad_op_func);

使之生效。

之后的代码:

const auto& dx_get_func = [&ctx, &grad_op_name]() -> const std::string& {

return ctx->GetOp(grad_op_name).output("dx", 0);

};

ctx->FwOp().InputGradBind(user_op::OpArg("x", 0), dx_get_func);

是将前向的输入 x 和刚刚设置的求梯度的方法的输出(dx) 绑定到一起,这样,使用 OneFlow 训练时,就可以自动求导。

其它

  • 本文涉及的代码
  • Op 注册的更多高级设置
  • 注册反向梯度时,也可以使用已有的 Op,而无需专门定制反向 Op。