关于深度学习:如何在OneFlow中新增算子

撰文|姚迟、郑泽康

本文将以开发一个 leaky_relu（精确说是 leaky_relu_yzh op，因为 master 分支的 leaky_relu 组合了其它知识点）为例介绍如何在 OneFlow 中新增算子（https://github.com/Oneflow-In…）。

1

背景

op 与 kernel

op 与 kernel 是两个有关联的概念。op 是逻辑上的算子，蕴含 OneFlow Compiler 在构建计算图时所须要的必要信息，如输出、输入形态，哪些张量须要主动求导等信息。有了 op 中的信息，OneFlow Compiler 就能够构建计算图并根据计算图做资源申请、构建等操作（如依据张量的输入输出大小申请内存），然而 op 中不蕴含具体的解决数据的逻辑。

在真正须要解决数据时，OneFlow Runtime 会启动 kernel 实现计算，所以 kernel 中蕴含了具体解决数据的逻辑。对于一个逻辑上的 op，OneFlow Runtime 会依据数据类型、硬件设施（比方是 CPU 还是 CUDA）的具体情况，抉择启动不同的 kernel。

OneFlow 中的零碎 op 与 user op

在 OneFlow 零碎中存在两类算子（op）：零碎 op 和 user op。

零碎 op 定义在：oneflow/core/operator/ 目录，对应的 kernel 实现在：oneflow/core/kernel 目录。零碎 op 是对构图、流水等零碎性能较为要害的一些 op。

除极少数 op 属于零碎 op 外，大多数 op 都是 user op，这些 user op 和用户模型业务逻辑相干。OneFlow user op 的定义及 kernel 实现别离在 oneflow/user/ops 和 oneflow/user/kernels 目录下。

目前 OneFlow 已实现了丰盛的算子库，然而当已有的算子库无奈满足搭建模型的需要时，就须要新增算子。本文介绍的新增算子指的是新增 user op。

ODS 与 TableGen

TableGen（https://llvm.org/docs/TableGe…）是一个代码生成工具，简略而言，它读取并解析一个 .td 格局（语法靠近 C++ 模板）的文件，而后交给 TableGen 后端
（https://llvm.org/docs/TableGe…）生成另外格局的语言。

MLIR 基于 TableGen 制订了一套算子定义标准 ODS（https://mlir.llvm.org/docs/Op…）以及对应的后端 OpDefinitionsGen（https://github.com/llvm/llvm-…。）

OneFlow 在 ODS 的根底上，实现了 TableGen OneFlow 后端（https://github.com/Oneflow-In…），并应用它来定义 OneFlow user op。

因而，OneFlow 的 user op 定义写在 OneFlowUserOps.td 文件中。

2

开发 op

在 OneFlow 中开发一个新的 user op，次要分为以下 4 步：

1. 定义 op
2. 实现 kernel 计算逻辑
3. 导出 functional 接口
4. 实现用于求导的反向逻辑

定义 op

定义 op 指的是，对 op 的名称，op 的输出、输入数据类型和 op 的属性进行申明。OneFlow 遵循 MLIR 的 ODS（Operation Definition Specification）（https://mlir.llvm.org/docs/Op…）实现了本人的 MLIR OneFlow Dialect。在算子定义方面，这样做的益处是，各种推导函数和序列化 / 反序列化的接口都能够委托给 ODS，升高了人工手写出错的概率，后续优化、格局转化等流程能够更灵便。

定义一个 OneFlow user op，次要包含 5 个局部，别离是：

• op class
• 输出 input
• 输入 output
• 属性 attrs
• 导出并实现推导接口

op class

能够在 oneflow/ir/include/OneFlow/OneFlowUserOps.td 查看 op 定义的源码。

以 def 关键字结尾定义一个 op，该 op 继承 OneFlow_BaseOp，同时指定 OneFlow_BaseOp 的模版参数。模版参数顺次为 op type name、Trait（https://mlir.llvm.org/docs/Tr…）列表。

def OneFlow_LeakyReluYZHOp : OneFlow_BaseOp<"leaky_relu_yzh", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {//...}

其中 “leaky_relu_yzh” 是指定的 op type name。每个 op 都须要指定一个全局惟一的 op type name 作为全局标识符。

第二个模板参数是一个 list（[...]），其中的每一项都是一个 Trait，OneFlow 中罕用的有：

• NoSideEffect 示意该算子无副作用（即不会扭转内存、网络、管道、磁盘等的零碎状态），这个个性能够领导某些优化操作
• NoGrad 示意该算子在数学上没有梯度（不可导）
• CpuOnly 示意该算子只反对在 CPU 设施上执行
• SupportNonContiguous 示意该算子是否反对 NonContiguous 张量（对于 Contiguous Tensor 的概念，能够参考 PyTorch Internals 中的相干内容）

输出 input 与输入 output

通过重写 input 域来定义 op 的输出，比方

// 一个输出 x
let input = (ins
  OneFlow_Tensor:$x
);

定义了一个输出张量 x。输出的格局为 输出类型:$name。

输出类型目前包含：

• OneFlow_Tensor
• Variadic<OneFlow_Tensor>：指可变 tensor，比方 concat op，反对 concat 可变个数的 tensor。
• Optional<OneFlow_Tensor>：示意这个 tensor 是可选的，既能够有也能够没有，比方 conv op 中的 add_output。

一个 op 也能够定义多个输出，比方：

// 两个输出：a, b
 let input = (ins
   OneFlow_Tensor:$a,
   OneFlow_Tensor:$b
 );

通过重写 output 域来定义 op 的输入，比方上面定义了 2 个输入张量：

let output = (outs
  OneFlow_Tensor:$out0,
  OneFlow_Tensor:$out1
);

属性 attrs

通过重写 attrs 域定义 op 的属性，比方定义 dropout（https://oneflow.readthedocs.i…）中的 rate 属性：

  let attrs = (ins
    DefaultValuedAttr<F32Attr, "0.">:$rate
  );

它示意名为 $rate 的类型是 F32Attr，默认值是 0.。这里也能够不指定默认值：

  let attrs = (ins
    F32Attr:$rate
  );

I32Attr、F32Attr、BoolAttr、StrAttr、I32ArrayAttr 等常见根底数据类型定义在 OpBase.td
（https://github.com/llvm/llvm-…）中。

OneFlow 自定义数据类型，如 ShapeAttr、DTArrayAttr 等定义在 OneFlowBase.td
（https://github.com/Oneflow-In…）中。

导出并实现推导接口

还有一些其它域，用于指定是否生成对应的接口。这些接口往往是构建计算图过程中的推导接口。

比方 shape 推导（依据输出的 shape 推导输入的推导）、data type 推导、SBP 推导等。

OneFlow-TableGen 仅负责生成这些函数的接口，开发者须要在其主动生成的 cpp 文件中实现这些接口。默认状况不会生成下列任何接口，开发者须要显式指定须要生成哪些接口。

  let has_check_fn = 1;                         // 生成属性查看接口
  let has_logical_tensor_desc_infer_fn = 1;     // 生成 logical shape 推导接口
  let has_physical_tensor_desc_infer_fn = 1;    // 生成 physical shape 推导接口
  let has_get_sbp_fn = 1;                       // 生成 get sbp 接口
  let has_sbp_signature_infer_fn = 1;           // 生成 sbp signature 推导接口，将来会移除，举荐应用 has_nd_sbp_infer_fn
  let has_data_type_infer_fn = 1;               // 生成 data type 推导接口
  let has_device_and_stream_infer_fn = 1;       // 生成 device 推导接口
  let has_input_arg_modify_fn = 1;              // 生成输出 modify 接口，比方设置 is_mutable、requires_grad（用于 Lazy）等
  let has_output_arg_modify_fn = 1;             // 生成输入 modify 接口，比方设置 is_mutable、requires_grad（用于 Lazy）等
  let has_output_blob_time_shape_infer_fn = 1;  // 生成输入 time shape 推导接口
  let has_nd_sbp_infer_fn = 1;                  // 生成 nd sbp 推导接口

个别罕用的是上面几个：

  let has_logical_tensor_desc_infer_fn = 1;
  let has_physical_tensor_desc_infer_fn = 1;
  let has_data_type_infer_fn = 1;
  let has_get_sbp_fn = 1;

理解完下面这些概念和用法后，能够开始批改 oneflow/ir/include/OneFlow/OneFlowUserOps.td文件。

leaky_relu_yzh op 残缺的定义见 这里（https://github.com/Oneflow-In…）。

在 OneFlowUserOps.td 中新增 Op 定义之后，从新 make 后会主动在 build 目录下的 oneflow/core/framework/ 目录下生成文件以下几个文件：

op_generated.h：由解析 .td 文件生成的 op C++ 类
op_generated.cpp：由解析 .td 文件生成的 op 注册代码（蕴含调用 REGISTER_USER_OP 宏的代码）

之后须要做的就是在 oneflow/user/ops（https://github.com/Oneflow-In…）目录下新加一个 cpp 文件，用于实现 op 的接口。

比方 leaky_relu_yzh 对应的文在 oneflow/user/ops/leaky_relu_yzh_op.cpp（https://github.com/Oneflow-In…），实现了推导逻辑张量、推导物理张量、推导 SBP 信息以及推导输入数据类型各接口。

实现 Kernel 逻辑

op 的计算反对多种设施（如 CPU、GPU、DCU 等），所以要别离实现计算逻辑。

相干代码：

• Leaky ReLU CPU Kernel
  （https://github.com/Oneflow-In…）
• Leaky ReLU GPU KernelCPU
  （https://github.com/Oneflow-In…）

计算逻辑

template<typename T>
class CpuLeakyReluYZHKernel final : public user_op::OpKernel {
 public:
  CpuLeakyReluYZHKernel() = default;
  ~CpuLeakyReluYZHKernel() = default;

 private:
  void Compute(user_op::KernelComputeContext* ctx) const override {const user_op::Tensor* x = ctx->Tensor4ArgNameAndIndex("x", 0);
    user_op::Tensor* y = ctx->Tensor4ArgNameAndIndex("y", 0);
    const int32_t elem_cnt = x->shape().elem_cnt();
    const T* x_ptr = x->dptr<T>();
    T* y_ptr = y->mut_dptr<T>();
    const auto alpha = ctx->Attr<float>("alpha");
    FOR_RANGE(int32_t, i, 0, elem_cnt) {y_ptr[i] = x_ptr[i] > 0 ? x_ptr[i] : alpha * x_ptr[i]; }
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false;}
};

在 OneFlow 中实现 kernel，必须定义一个继承自 oneflow::user_op::OpKernel 的类，并重写其中的虚函数。在以上代码中，重写了 Compute 与 AlwaysComputeWhenAllOutputsEmpty 两个虚函数，它们的意义别离是：

• Compute 必须重写，在其中实现具体的运算逻辑
• AlwaysComputeWhenAllOutputsEmpty 必须重写，对于绝大多数 op 而言，间接返回 false 即可。对于极少数外部须要保护状态，即便输入为空也须要调用 kernel 进行计算的 op 而言，应该返回 true

Compute 办法中通过调用 user_op::KernelComputeContext* ctx 中的接口，能够获取输出张量、输入张量、attr 具体的数据，再依照算子的算法逻辑对它们进行解决。以下是对 CpuLeakyReluKernel::Compute 解决逻辑的解读：

• 首先获得 “x“，”y” 两个 Tensor。传入 Tensor4ArgNameAndIndex 的字符串要和之前在 OneFlowUserOps.td 设置的名称统一
• 获取 x 的元素个数，以便后续用于 for 循环进行计算
• 获取属性 alpha
• 进入次数为 elem_cnt 的 for 循环，将后果写入

注册 Kernel

实现 kernel 类后，须要调用 REGISTER_USER_KERNEL 注册。

#define REGISTER_CPU_LEAKY_RELU_YZH_KERNEL(dtype)                     \
  REGISTER_USER_KERNEL("leaky_relu_yzh")                              \
      .SetCreateFn<CpuLeakyReluYZHKernel<dtype>>()                    \
      .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCPU) \
                       && (user_op::HobDataType("y", 0) == GetDataType<dtype>::value));

这里会调用 REGISTER_USER_KERNEL 宏，包含以下信息：

1. op type name：为哪个 op 注册 kernel
2. SetCreateFn<T>()：该模板办法的模板参数 T，就是咱们实现的 kernel 类，OneFlow Runtime 将应用它创立 kernel 对象。
3. SetIsMatchedHob：因为一个 op 可能有多个 kernel，要想依据物理设施及数据格式的不同而抉择不同的 kernel 进行计算，就须要调用 SetIsMatchedHob 进行设置。该办法承受一个表达式，表达式为 true 时，OneFlow 将调用该 kernel 实现计算。以上代码的匹配逻辑是：当硬件设施为 cpu，且 y 的数据类型和 dtype 统一时，抉择调用注册的 kernel 类（CpuLeakyReluYZHKernel<dtype>）。

GPU 计算逻辑

CUDA 编程基础知识入门能够参考：

• 视频：CUDA 的由来（https://www.bilibili.com/vide…）
• 视频：CUDA 的入门小程序（https://www.bilibili.com/vide…）
• 视频：线程层级（https://www.bilibili.com/vide…）

不过以上的视频都无奈代替本人认真学习官网材料：CUDA C Programming Guide（https://docs.nvidia.com/cuda/…）

理解了 CUDA 的基础知识，就不难理解 leaky_relu CUDA 版本的实现。
首先定义了 leaky_relu 前向运算的 CUDA 核函数

template<typename T>
__global__ void LeakyReluForwardGpu(const int n, const float alpha, const T* x, T* y) {CUDA_1D_KERNEL_LOOP(i, n) {y[i] = x[i] > 0 ? x[i] : x[i] * alpha; }
}

其中调用了宏 CUDA_1D_KERNEL_LOOP（https://github.com/Oneflow-In…）进行运算

在 Compute 函数中，调用了 RUN_CUDA_KERNEL (也是定义在 cuda_util.h 这个文件中)这个宏启动核函数。

对应的 GPU kernel 类的实现见：

https://github.com/Oneflow-In…

其中用到了启动 kernel 的宏 RUN_CUDA_KERNEL，它的定义是:

#define RUN_CUDA_KERNEL(func, device_ctx_ptr, thread_num, ...)           \
  func<<<SMBlocksNum4ThreadsNum(thread_num), kCudaThreadsNumPerBlock, 0, \
         (device_ctx_ptr)->cuda_stream()>>>(__VA_ARGS__)

1. 第一个参数是核函数名字
2. 第二个参数是 device context，后续获取对应的 cuda_stream
3. 第三个参数是要启动的线程数量，会依据线程数量来计算所需的 Block 数目。
   因为 leaky relu 是 elementwise 运算，各个元素互不影响，所以咱们启动了 elem_cnt 个线程。

后续的注册与 CPU 版本相似，这里不再赘述。间接参考以下代码即可：

https://github.com/Oneflow-In…

能够看到不同设施类的 Compute 中大部分代码是反复的。一种更优的代码组织形式是用一个 .cpp 文件实现 kernel 和注册的逻辑，.cu 文件编写 GPU Kernel 函数和 GPU 模板特化的代码，.h 文件用于定义和编写注册宏。可参考 dim_gather_kernel_*
（https://github.com/Oneflow-In…）中的代码。

OneFlow 为了适配多种设施，还提供了 Primitive 组件，能够参考：Primitive PR（https://github.com/Oneflow-In…）

导出 functional 接口

对于 functional 接口层的具体介绍在这里：https://github.com/Oneflow-In…

概括而言，functional 层起到了“上接 Python，下联 C++”的作用：

   ┌─────────────┐
   │   Module    │
   │  (Python)   │
   ├─────────────┤
   │             │
   │ Functional  │
   ├─────────────┤
   │             │
   │ Op/Kernels  │
   │   (C++)     │
   └─────────────┘

因而，在上文定义 op 和注册 kernel 后，须要为算子导出 functional 接口，能力使用户通过 Python 代码调用该算子。

导出 functional 接口分为以下几个步骤：

1. 实现对应的 functor 并注册
2. 在 oneflow/core/functional/functional_api.yaml 中增加接口形容

实现对应的 functor 并注册

对于 leaky_relu_yzh op，在 activation_functor.cpp
（https://github.com/Oneflow-In…）中，对其进行定义：

class LeakyReluYZHFunctor {
 public:
  LeakyReluYZHFunctor() {op_ = CHECK_JUST(one::OpBuilder("leaky_relu_yzh").Input("x").Output("y").Build());
  }
  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const float& alpha) const {
    MutableAttrMap attrs;
    JUST(attrs.SetAttr<float>("alpha", alpha));
    return OpInterpUtil::Dispatch<one::Tensor>(*op_, {x}, attrs);
  }

 private:
  std::shared_ptr<OpExpr> op_;
};

• 在构造函数里，结构了 leaky_relu 这个 op
• 实现 operator() 重载运算符，通过 Dispatch 调用结构好的 op，并别离传入输出，属性

相似的咱们也给 LeakyReluGrad 导出 functional 接口，以便后续编写求导逻辑应用。

最初咱们须要注册到 Functional Library：

https://github.com/Oneflow-In…

m.add_functor<impl::LeakyReluYZHFunctor>("LeakyReluYZH"); // 留神最初字符串中的名字在后续的 functional_api.yaml 中会用到

通过 m.add_functor 注册后的 functor，能够在 C++ 层应用，如通过 functional::LeakyRelu 就能够调用 LeakyReluFunctor。

在 functional_api.yaml 中增加接口形容

functional 通过解析 yaml 配置文件，在 build 过程中主动帮咱们生成接口。
在 functional_api.yaml
（https://github.com/Oneflow-In…）文件中，编写相干配置。

*https://github.com/Oneflow-In…

 - name: "leaky_relu_yzh"
  signature: "Tensor (Tensor x, Float alpha) => LeakyReluYZH"
  bind_python: True

• 其中 name 示意导出到 Python 接口后函数的名字，比方导出后在 Python 下应用就是

• signature 用于形容接口原型及 C++ 代码的对应关系。=> 右边的为原型；=> 左边为对应的 Functional Library 中的名字。这里LeakyRelu 和后面导出时指定的字符串是统一的。
• bind_python，示意这个接口是否须要绑定 Python 接口。比方上面的 leaky_relu_grad，咱们不会在 Python 层用到（但会在 C++ 层求导应用），所以设置为 False。

实现以上工作后，新增的算子曾经反对正向运算，编译好代码便能够进行如下简略的测试：

import oneflow as flow 
import numpy as np


x_tensor = flow.Tensor(np.random.randn(3, 3))
out = flow._C.leaky_relu_yzh(x_tensor, alpha=0.2)

然而，还须要注册反向，能力反对反向流传。咱们也先将反向须要的 LeakyReluGrad 导出为 functional 接口。

- name: "leaky_relu_yzh_grad"
  signature: "Tensor (Tensor x, Tensor dy, Float alpha) => LeakyReluYZHGrad"
  bind_python: False 实现用于求导的反向逻辑

反向流传的实质就是高数中的链式法则，只不过 Autodiff 将链式法则变得模块化、易复用。

能够先浏览 CSC321 Lecture 10: Automatic Differentiation（https://www.cs.toronto.edu/~r…）理解 autodiff 的基本概念。

从逻辑上而言，一个算子在反向过程中可能求导数，个别须要以下信息：

• 正向过程中的输出、输入
• 正向过程的 attr
• 反向过程中上一层（正向过程中的下一层）传递过去的正向输入的梯度

将来 Graph 模式和 Eager 模式下的反向逻辑会合并，但目前还是须要别离注册。

为 Eager 模式注册反向

求导局部在 oneflow/core/autograd/gradient_funcs/activation.cpp
（https://github.com/Oneflow-In…）实现

次要有以下几局部：

• LeakyReluYZHCaptureState：用于存储数据的构造体

这是一个简略的构造体，继承自 AutoGradCaptureState，用于存储 op 的属性，以便于后续求导。

struct LeakyReluYZHCaptureState  : public AutoGradCaptureState {
  bool requires_grad; // 输出 x 是否须要梯度
  float alpha=0.0; // 输出的参数 alpha
};

• LeakyReluYZH 类：继承自 OpExprGradFunction 的类。须要重写三个函数：Init、Capture、Apply。

class LeakyReluYZH : public OpExprGradFunction<LeakyReluYZHCaptureState> {
 public:
  Maybe<void> Init(const OpExpr& op) override {//...}

  Maybe<void> Capture(LeakyReluYZHCaptureState* ctx, const TensorTuple& inputs,
                      const TensorTuple& outputs, const AttrMap& attrs) const override {//...}

  Maybe<void> Apply(const LeakyReluYZHCaptureState* ctx, const TensorTuple& out_grads,
                    TensorTuple* in_grads) const override {//...}
};

• Init：做的是一些初始化的工作，能够依据前向 op 的配置，来初始化属性。

  Maybe<void> Init(const OpExpr& op) override {const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
   CHECK_NOTNULL_OR_RETURN(fw_op_expr);
   base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
   return Maybe<void>::Ok();}
  

• Capture：用于捕获前向的 Tensor，属性，用于后续的求导。

以 LeakyReluYZH 为例子，咱们须要失去：a) 输出的 Tensor，当 Tensor 数值大于 0，梯度为 1，当小于 0，梯度为 alpha b) alpha 的数值

Maybe<void> Capture(LeakyReluYZHCaptureState* ctx, const TensorTuple& inputs,
                    const TensorTuple& outputs, const AttrMap& attrs) const override {CHECK_EQ_OR_RETURN(inputs.size(), 1);                      // 判断输出个数是否为 1
  ctx->requires_grad = inputs.at(0)->requires_grad();        // 判断输出是否须要梯度
  if (!ctx->requires_grad) {return Maybe<void>::Ok(); }     // 如果不须要梯度，也就不须要求导了，间接返回 Maybe<void>::Ok()

  ComposedAttrMap composed_attrs(attrs, base_attrs_);
  ctx->alpha = JUST(composed_attrs.GetAttr<float>("alpha")); // 获取 alpha，并存入 ctx->alpha 中
  ctx->SaveTensorForBackward(inputs.at(0));                  // 调用 SaveTensorForBackward 办法，保留输出的 Tensor
  return Maybe<void>::Ok();}

• Apply：理论计算梯度的函数，咱们能够拿到先前的 Tensor，并调用 functional 接口下注册的 LeakyReluGrad，求得梯度，并返回

  Maybe<void> Apply(const LeakyReluYZHCaptureState* ctx, const TensorTuple& out_grads,
                    TensorTuple* in_grads) const override {CHECK_EQ_OR_RETURN(out_grads.size(), 1);  // 查看梯度 Tensor 个数是否为 1
    in_grads->resize(1);                      // 因为输出只有一个，所以咱们 resize(1)
    if (ctx->requires_grad) {const auto& x = ctx->SavedTensors().at(0); // 调用 SavedTensors 接口，拿到先前通过 SaveTensorForBackward 接口保留的 Tensor
      in_grads->at(0) = JUST(functional::LeakyReluYZHGrad(x, out_grads.at(0), ctx->alpha)); // 拿到 x，dy，alpha，传入给 LeakyReluYZHGrad 计算，并将梯度返回给 in_grads->at(0)
    }
    return Maybe<void>::Ok();}

最初一步是注册，第一个参数是 op type name，第二个参数是继承自 OpExprGradFunction 的类。

REGISTER_OP_EXPR_GRAD_FUNCTION("leaky_relu_yzh", LeakyReluYZH); // 第二个参数是用于求导的类名

为 Graph 模式注册反向

为 Graph 模式注册 leaky_relu_yzh op 的反向代码在：

https://github.com/Oneflow-In…

REGISTER_USER_OP_GRAD("leaky_relu_yzh")
    .SetBackwardOpConfGenFn([](user_op::BackwardOpConfContext* ctx) -> Maybe<void> {
      // 依据前向的 op type name，拼凑出一个 leaky_relu_yzh_grad_op_name（leaky_relu_yzh_grad）const std::string leaky_relu_yzh_grad_op_name = ctx->FwOp().op_name() + "_grad";
      
      ctx->DefineOp(leaky_relu_yzh_grad_op_name, [&ctx](user_op::BackwardOpBuilder& builder) {
        // 构建一个 op（op type name 为 leaky_relu_yzh_grad）// 把前向输入 y 的梯度，作为 leaky_relu_yzh_grad 的输出 dy
        // 把前向的 x 作为 leaky_relu_yzh_grad 的输出 x
        // 输入为 dx
        // attr alpha 同前向一样        
        return builder.OpTypeName("leaky_relu_yzh_grad")
            .InputBind("dy", ctx->FwOp().output_grad("y", 0))
            .InputBind("x", ctx->FwOp().input("x", 0))
            .Attr("alpha", ctx->FwOp().attr<float>("alpha"))
            .Output("dx")
            .Build();});

      // 把 leaky_relu_yzh_grad_op_name 算子的输入 dx 的后果
      // 绑定到前向输出 x 的反向梯度上
      // 即：//  leaky_relu_yzh 的输出 x 的梯度 = leaky_relu_yzh_grad 的输入 dx
      ctx->FwOp().InputGradBind(user_op::OpArg("x", 0),
                                [&ctx, &leaky_relu_yzh_grad_op_name]() -> const std::string& {return ctx->GetOp(leaky_relu_yzh_grad_op_name).output("dx", 0);
                                });
      return Maybe<void>::Ok();});

3

测试与文档

本文笼罩的内容实现后，只是“跑通”算子，还须要进一步欠缺，包含为算子增加测试和 API 文档，这些将在后续的文章中介绍。

欢送下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In…