组网只会拼装API？这套自定义算子教程让你的组网更加灵活多变

发布日期：2021-12-04T09:36:13.000+0000 浏览量：1245次

完整项目已在AI Studio开源，点击链接即可运行：

https://aistudio.baidu.com/aistudio/projectdetail/2326026

算子（Operator，简称Op）是构建神经网络的基础组件。在网络模型中，算子对应层中的计算逻辑，例如：卷积层（Convolution Layer）是一个算子；全连接层（Fully-connected Layer， FC layer）中的权值求和过程，是一个算子。学会定制化算子的C++实现可以更深入地了解神经网络运行的底层逻辑。

图1：算子(Operator)也简称为OP，可以理解为一个计算函数

构建神经网络模型需要使用到各种各样的算子（Operator），例如卷积算子、各类数学运算算子、激活函数算子等，深度学习框架需要提供这些算子的运算支持，因此在深度学习框架内部很大一部分工作都是围绕如何对这些算子（Oparetor）进行管理、调度和执行来展开。飞桨开源框架目前已经提供了大量的算子供用户使用，（如图2所示）在框架内部使用OpInfoMap对所有飞桨支持的算子（Oparetor）进行注册和管理（OpInfoMap本质上是一个全局Map，使用算子的名称作为Key，算子的描述信息、构建函数、辅助行功能函数等具体组件作为对应的Value项）。

算子的运算逻辑如果按编程语言来分，可以分为Python层和C++层，如图2所示。我们平常使用的算子，其底层逻辑是用C++实现的，通过一些封装后，便以Python接口的形式展现给开发者，而两者之间的互操作是通过Pybind实现的。算子的底层逻辑中，TraceOp是飞桨动态图模式下进行模型训练最关键的组件，简单来说，在Python层调用算子（Oparetor）执行运算操作时，通过TraceOp在OpInfoMap 找到对应的算子（Oparetor）信息，构建算子，随后找到其计算kernel并调用，完成计算并返回结果。

这里可以看出算子（Oparetor）和Kernel是两个不同的概念，算子是一个运算逻辑的表示（如加减乘除运算），Kernel是运算逻辑真正进行计算处理的实现（需要分配具体的计算设备完成计算，如CPU Kernel、GPU Kernel等）

换句话说，因为硬件的不同，相同的算子需要不同的Kernel，比如你实现了一个在CPU上执行的计算Kernel，那么这个算子可以在CPU上运行，要想在别的计算平台上运行，还需要实现该平台的Kernel。

图2：PaddlePaddle Op算子体系(动态图模式)

上图简单介绍了算子的执行逻辑，如果要更进一步地解释算子的执行逻辑，需要用一个例子来说明。当我们在动态图模式下执行Y=relu(X)时，框架会通过代理对象TraceOp到OpInfoMap里找被调用的算子，比如调用ReluOp，它会找到该算子并调用其计算Kernel，然后完成计算及返回结果，具体的计算过程可以分成2个步骤：

调用relu算子的forward计算函数完成Y的计算。
创建backward所需的Op算子以及输入输出变量(此时不进行计算，待后续调用backward()后才会进行反向计算)。

图3：Op算子前向、反向计算(动态图模式)

这里需要补充说明的是，神经网络的“学习”，其实是通过调整权重的方式来实现的，当神经网络的输出与期望输出的误差趋近于0时，“学习”的过程也就完成了。算子作为一个神经网络的基本组成部分，其前反向计算的结果有2个作用：

1.通过前向计算得到输出结果。

2.使用前向计算结果与期望结果的差值进行反向梯度的计算，并利用反向传播计算出的梯度值更新权重。

因此，前反向计算的结果将会影响模型最后的训练结果，如果反向传播的计算公式推导错误，模型将无法找到最优或近似最优解。

总的来说，大部分的自定义算子其代码在格式上是固定的，写法也都是类似的，基本上仿照其他算子的格式就可以完成；但每个算子前向和反向的计算逻辑都是不同的。在实现时除了要正确实现计算公式之外，还要尽可能的保证计算的性能，这才是实现算子的难点所在。

C++自定义算子格式

基本格式

在编写运算函数之前，需要引入 PaddlePaddle 扩展头文件：

#include "paddle/extension.h"

算子运算函数有特定的函数写法要求，在编码过程中需要遵守，基本形式如下：

std::vector<paddle::Tensor> OpFucntion(const paddle::Tensor& x, ..., int attr, ...) {
  ...
}

这一部分其实就是固定格式，所有用C++编写的Paddle算子都需要使用这个格式。换句话说，这是Paddle提供的算子接口，只需要按照这个接口定义算子即可。

适配多种数据类型

在实际开发中，一个算子往往需要支持多种数据类型，这时我们可以将计算函数包装成 模板函数 ，通过 switch-case 语句实现支持多种数据类型的操作：

switch(x.type()) {
  case paddle::DataType::FLOAT32:
    ...
    break;
  case paddle::DataType::FLOAT64:
    ...
    break;
  default:
    PD_THROW(
      "function  ... is not implemented for data type `",
      paddle::ToString(x.type()), "`");
}

如果不想使用 switch-case 来实现，也可以使用官方提供的DISPATCH宏，如PD_DISPATCH_FLOATING_TYPES，使用方式可参考后面给出的代码示例。

维度与类型的推导

飞桨同时支持动态图与静态图的执行模式，在静态图模式下，组网阶段需要完成 Tensor shape 和 dtype 的推导，从而生成正确的模型描述，用于后续Graph优化与执行。因此，除了算子的运算函数之外，还需要实现前向运算的维度和类型的推导函数。

维度推导（InferShape）和类型推导（InferDtype）的函数写法也是有要求的，格式如下：

需要注意的是，输入输出参数与Forward计算函数的输入输出Tensor应该按顺序一一对应：

对于仅有一个输入Tensor和一个输出Tensor的自定义算子，如果输出Tensor和输入Tensor的shape和dtype一致,可以省略InferShape和InferDtype函数的实现，但其他场景下均需要实现这两个函数。

自定义算子注册

最后，需要调用 PD_BUILD_OP 系列宏，构建算子的描述信息，并关联前述算子运算函数和维度、类型推导函数。其格式如下：

需要注意的是：

PD_BUILD_OP 用于构建前向算子，其括号内为算子名，也是后面在python端使用的接口名，注意前后不需要引号，注意该算子名不能与 PaddlePaddle 内已有算子名重名
PD_BUILD_GRAD_OP 用于构建前向算子对应的反向算子，PD_BUILD_DOUBLE_GRAD_OP 用于构建前反向算子对应的二次求导算子。Paddle目前支持的多阶导数只支持到二阶导

动手实现

CPU算子

下面将以一个比较简单的sin函数为例，自定义一个CPU算子（本文主要对CPU算子进行讲解，GPU算子可参考：飞桨官方文档-自定义外部算子教程）。

导入必要的头文件

#include "paddle/extension.h"
#include <vector>
#define CHECK_CPU_INPUT(x) PD_CHECK(x.place() == paddle::PlaceType::kCPU, #x " must be a CPU Tensor.")

引入 PaddlePaddle 扩展头文件以及宏定义，检验输入的格式。

实现forward计算函数

为了适配多种数据类型，这里可以在计算函数中引入模板。

前向计算最重要的就是实现计算函数，C++里提供了一些基础运算的函数，可以直接使用，基本语法一般为

std::function(input)。
template <typename data_t> // 模板函数
void sin_cpu_forward_kernel(const data_t* x_data,
                            data_t* out_data,
                            int64_t x_numel) {
  for (int i = 0; i < x_numel; ++i) {
    out_data[i] = std::sin(x_data[i]);
  }
}

接着只需要将前面实现的计算函数按照前面给的格式套进前向传播即可：

std::vector<paddle::Tensor> sin_cpu_forward(const paddle::Tensor& x) {
  CHECK_CPU_INPUT(x);
// 声明输出变量out，需传入两个参数（运行的设备类型及维度信息）
  auto out = paddle::Tensor(paddle::PlaceType::kCPU, x.shape());

  // 计算实现，PD_DISPATCH宏具体使用方法和原理详见:
// 飞桨官方文档自定义外部算子CPU实现部分 
  PD_DISPATCH_FLOATING_TYPES(
      x.type(), "sin_cpu_forward_kernel", ([&] {
        sin_cpu_forward_kernel<data_t>( // 调用前面定义好的前向计算函数，这里必须是data_t
            x.data<data_t>(), // 获取输入的内存地址，即从内存空间中取数据，这里必须是data_t
            out.mutable_data<data_t>(x.place()), x.size()); // 为输出申请内存空间，这里必须是data_t
      }));

  return {out};
}

实现backward计算函数

这部分需要一定的数学基础，要了解 偏微分 的计算方法，理解神经网络的梯度概念，我在实现过程中也查阅了一些资料，给大家分享：

3blue1brown： https://www.3blue1brown.com/lessons/backpropagation-calculus
神经网络之梯度下降法及其实现：

https://blog.csdn.net/nanhuaibeian/article/details/100184893
wolframalpha： https://www.wolframalpha.com/

最后一个网站是一个可以直接计算偏导数的网站，比较方便，比如这里需要计算sin函数的偏导:

反向传播最难的就是计算梯度，如果会计算，其实就很简单了，跟前向计算是类似的：

template <typename data_t>
void sin_cpu_backward_kernel(const data_t* grad_out_data,
                              const data_t* x_data,
                              data_t* grad_x_data,
                              int64_t out_numel) {
  for (int i = 0; i < out_numel; ++i) {
    grad_x_data[i] = grad_out_data[i] * std::cos(x_data[i]); // 结果是返回的梯度值乘函数导数值
  }
}

std::vector<paddle::Tensor> sin_cpu_backward(const paddle::Tensor& x, // forward的输入
                                              const paddle::Tensor& out, // forward的输出
                                              const paddle::Tensor& grad_out) { // backward的梯度变量

  auto grad_x = paddle::Tensor(paddle::PlaceType::kCPU, x.shape());

  // 计算实现
  PD_DISPATCH_FLOATING_TYPES(out.type(), "sin_cpu_backward_kernel", ([&] {
                               sin_cpu_backward_kernel<data_t>(
                                   grad_out.data<data_t>(), // 获取内存地址，即从内存空间中取数据
                                   x.data<data_t>(), // 获取内存地址，即从内存空间中取数据
                                   grad_x.mutable_data<data_t>(x.place()), // 申请内存空间
                                   out.size()); // 传入输出的维度信息
                             }));

  return {grad_x};
}

维度推导

维度推导部分其实只需要根据格式实现InferShape和InferDtype函数即可：

// 维度推导
std::vector<std::vector<int64_t>> sinInferShape(std::vector<int64_t> x_shape) {
  return {x_shape};
}

// 类型推导
std::vector<paddle::DataType> sinInferDtype(paddle::DataType x_dtype) {
  return {x_dtype};
}

因为sin(x)函数输入和输出的维度一致，所以可以省略InferShape和InferDtype函数的实现。

自定义算子注册

最后也是按照格式完成自定义算子的注册即可：

PD_BUILD_OP(custom_sin_cpu)
    .Inputs({"X"})
    .Outputs({"Out"})
    .SetKernelFn(PD_KERNEL(sin_cpu_forward))
    .SetInferShapeFn(PD_INFER_SHAPE(sinInferShape))
    .SetInferDtypeFn(PD_INFER_DTYPE(sinInferDtype));

PD_BUILD_GRAD_OP(custom_sin_cpu)
    .Inputs({"X", "Out", paddle::Grad("Out")})
    .Outputs({paddle::Grad("X")})
    .SetKernelFn(PD_KERNEL(sin_cpu_backward));

自定义

CPU算子的使用

使用JIT (即时编译)安装加载自定义算子，其基本格式如下：

为了便于读者体验，我已经将算子写好，并打包到了AI Studio的项目【飞桨高阶使用教程：自定义算子的实现和使用】中，位于custom_op/custom_sin_cpu.cc，直接调用即可：

from paddle.utils.cpp_extension import load
custom_ops = load(
    name="custom_jit_ops",
    sources=["custom_op/custom_sin_cpu.cc"])

custom_sin_cpu = custom_ops.custom_sin_cpu
Compiling user custom op, it will cost a few seconds.....
使用该算子也非常简单，直接使用即可，如下所示：
import paddle
import paddle.nn.functional as F
import numpy as np

# 定义执行环境
device = 'cpu'
paddle.set_device(device)

# 将输入数据转换为张量
data = np.random.random([4, 12]).astype(np.float32)
x = paddle.to_tensor(data, stop_gradient=False)

# 调用自定义算子实现前向计算
y = custom_sin_cpu(x)
# 调用自定义算子实现反向传播
y.mean().backward()

print("前向计算结果：{}".format(y))
print("梯度结果：{}".format(x.grad))

为了验证算子的正确性，我们可以跟Paddle现有的算子做对比，看看前向传播和梯度的计算结果是否一致：

import paddle
import paddle.nn.functional as F
import numpy as np

device = 'cpu'
paddle.set_device(device)
data = np.random.random([4, 12]).astype(np.float32)
x_target = paddle.to_tensor(data, stop_gradient=False)
y_target = paddle.sin(x_target)
y_target.mean().backward()
x = paddle.to_tensor(data, stop_gradient=False)
y = custom_sin_cpu(x)
y.mean().backward()

# 输出都为True表示结果正确
print("sin_result: ",paddle.allclose(y_target, y).numpy())
print("sin_grad_result: ",paddle.allclose(x_target.grad, x.grad).numpy())

sin_result:  [ True]
sin_grad_result:  [ True]

从输出结果可以看出，我们自定义的算子从实现功能上来说是正确的。

总结与升华

最后总结一下C++自定义算子最主要的思路，其实就是3点：

forward和backward实现
包装forward和backward函数并注册
编译加载并调用算子

从我的感受来说，我认为第一点是最为重要的部分，特别是反向传播里梯度的计算，需要一定的数学基础，要对神经网络的工作机制有较为深刻的理解。

本文仅介绍了如何实现一个在CPU上使用的算子，飞桨自定义算子机制也可以实现在GPU上使用，以及在CPU和GPU可以同时使用，具体可参考飞桨官方文档-自定义外部算子。

参考资料：

《飞桨论文复现打卡营》高阶课：Paddle自定义算子实现和使用

https://aistudio.baidu.com/aistudio/education/group/info/24681

论文复现营C++自定义算子开发Demo by zhangyunfei：

https://aistudio.baidu.com/aistudio/projectdetail/2296948

飞桨使用教程：自定义算子之自定义C++算子实现：

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/07_new_op/new_custom_op_cn.html

关注公众号，获取更多技术内容~

上一篇：嚯，这款AI建模工具实在太强大了，快来pick！

下一篇：Vision Transformer打卡营今晚正式结营！