撰文|姚迟、郑泽康
本文将以开发一个 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步:
- 定义 op
- 实现 kernel 计算逻辑
- 导出 functional 接口
- 实现用于求导的反向逻辑
定义 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宏,包括以下信息:
- op type name:为哪个 op 注册 kernel
SetCreateFn<T>()
:该模板方法的模板参数 T,就是我们实现的 kernel 类,OneFlow Runtime 将使用它创建 kernel 对象。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__)
- 第一个参数是核函数名字
- 第二个参数是 device context,后续获取对应的 cuda_stream
- 第三个参数是要启动的线程数量,会根据线程数量来计算所需的 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 接口分为以下几个步骤:
- 实现对应的 functor 并注册
- 在 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...
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用
。你还可以使用@
来通知其他用户。