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<std::string>("device_sub_tag", "py")
.Input("in")
.Output("out")
.SetTensorDescInferFn(
[](user_op::InferContext *ctx) -> Maybe<void> {
*ctx->Shape4ArgNameAndIndex("out", 0) =
*ctx->Shape4ArgNameAndIndex("in", 0);
*ctx->Dtype4ArgNameAndIndex("out", 0) =
*ctx->Dtype4ArgNameAndIndex("in", 0);
return Maybe<void>::Ok();
});
} // namespace
} // namespace oneflow
分析以上代码：
- oneflow/core/framework/framework.h 中包含了创建一个 op 所需要的所有接口
- .Attr<std::string>("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 的设置：
- Input("in") 表示其有一个名为 "in" 的输入
- Output("out") 表示其有一个名为 "out" 的输出
- 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<std::string>("device_sub_tag", "py")
.SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
const Shape* dy_shape = ctx->Shape4ArgNameAndIndex("dy", 0);
Shape* dx_shape = ctx->Shape4ArgNameAndIndex("dx", 0);
*dx_shape = *dy_shape;
return Maybe<void>::Ok();
});
值得注意的是，同前向类似，以上代码中 .Attr<std::string>("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。