跳转至

使用 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_nameOp 以及 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_forwardop_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_myreluUserOpConfBuilder 对象。

该对象包含 OpInput 等方法,用于封装自定义 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
  • 参数只有一个,类型为 tupletuple 中的元素个数和顺序,与 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.pyuser_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,可以参考 这里