飞桨高阶使用教程：自定义CPU算子的实现和使用_paadleinference cpu源码编译

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

张量与算子

• 张量(Tensor)，可以理解为多维数组,可以具有任意多的维度,不同Tensor可以有不同的数据类型(dtype)和形状(shape)。
• 算子(Operator)也简称为OP，负责对Tensor执行各种运算处理,可以理解为一个计算函数,函数的输入和输出都为Tensor。在PaddlePaddle中定义了大量的Operator来完成常见神经网络模型的Tensor运算处理,如conv2d, pool2d, Relu等

一、底层原理

PaddlePaddle中所有Op算子都会注册在OpInfoMap中,在Python端调用Op执行运算操作时,通过TraceOp在OpInfoMap找到对应的Op并调用其计算kernel ,完成计算并返回结果。

换句话说，因为硬件的不同，相同的算子需要不同的kernel，比如你写了一个在CPU上执行的算子，那么这个算子只能在CPU上运行，要想在别的计算平台上运行，还需要实现该平台的kernel。这篇文章讲的主要是怎么实现一个能在CPU上运行的算子。

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

例如，在动态图模式执行Y=relu(X)时,框架会通过TraceOp来完成:

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

PaddlePaddle Op算子正反向计算(动态图模式)

通过上面展示的底层原理，其实不难发现，一个算子最关键的部分就是前向传播与反向计算，这两个部分是算子的核心。

C++自定义算子内部原理(动态图模式)

在动态图模式执行Y=custom_ relu(X)时:使用C++自定义算子与原生算子的执行流程相同。

但和原生算子区别是原生算子随框架一起编译；自定义算子可单独编译。但最后都会注册到OpInfoMap中。

二、C++自定义算子格式

1.基本格式

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

#include "paddle/extension.h"
1

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

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

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

2.适配多种数据类型

在实际开发中，一个算子往往需要支持多种数据类型，这时就需要用到模板类。在上面接口上方定义：

template <typename data_t>
1

需要注意的是：模板参数名 data_t 用于适配不同的数据类型，不可更改为其他命名，否则会编译失败

然后通过 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()), "`");
}
1
2
3
4
5
6
7
8
9
10
11
12

如果不想使用 switch-case 来实现，也可以使用官方提供的dispatch宏，如PD_DISPATCH_FLOATING_TYPES

3.维度与类型的推导

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

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

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

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

4.自定义算子注册

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

PD_BUILD_OP(op_name)
    .Inputs({"X"})
    .Outputs({"Out"})
    .SetKernelFn(PD_KERNEL(Forward))
    .SetInferShapeFn(PD_INFER_SHAPE(ReluInferShape))
    .SetInferDtypeFn(PD_INFER_DTYPE(ReluInferDtype));

PD_BUILD_GRAD_OP(op_name)
    .Inputs({"X", "Out", paddle::Grad("Out")})
    .Outputs({paddle::Grad("X")})
    .SetKernelFn(PD_KERNEL(ReluCPBackward));
1
2
3
4
5
6
7
8
9
10
11

需要注意的是：

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

三、动手实现CPU算子

下面将以一个比较简单的Sin函数为例，自定义一个CPU算子。

1.导入必要的头文件

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

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

2.实现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]);
  }
}
1
2
3
4
5
6
7
8

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

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

  // 计算实现
  PD_DISPATCH_FLOATING_TYPES(
      x.type(), "sin_cpu_forward_kernel", ([&] {
        sin_cpu_forward_kernel<data_t>( // 调用前面定义好的前向计算函数
            x.data<data_t>(), // 获取输入的内存地址，即从内存空间中取数据
            out.mutable_data<data_t>(x.place()), x.size()); // 为输出申请内存空间 
      }));

  return {out};
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

3.实现backward计算函数

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

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

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

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

template <typename data_t>
void sin_cpu_backward_kernel(const data_t* grad_out_data,
                              const data_t* out_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(out_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>(), // 获取内存地址，即从内存空间中取数据
                                   out.data<data_t>(), // 获取内存地址，即从内存空间中取数据
                                   grad_x.mutable_data<data_t>(x.place()), // 申请内存空间
                                   out.size()); // 传入输出的维度信息
                             }));

  return {grad_x};
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

4.维度推导

维度推导部分其实只需要根据格式实现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};
}
1
2
3
4
5
6
7
8
9

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

5.自定义算子注册

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

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));
1
2
3
4
5
6
7
8
9
10
11

四、自定义CPU算子的使用

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

在本项目中，我已经将算子写好，位于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
1
2
3
4
5
6

Compiling user custom op, it will cost a few seconds.....
1

使用该算子也非常简单，直接使用即可，如下所示：

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))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

前向计算结果：Tensor(shape=[4, 12], dtype=float32, place=CPUPlace, stop_gradient=False,
       [[0.39883998, 0.46783343, 0.68504739, 0.44232291, 0.62964708, 0.71694267, 0.07653319, 0.51490635, 0.81647098, 0.68105453, 0.10945870, 0.08908488],
        [0.47923982, 0.01490644, 0.13291596, 0.21918269, 0.38499439, 0.75070190, 0.52795607, 0.05496189, 0.43035936, 0.17001969, 0.64533097, 0.22776006],
        [0.78449929, 0.54673332, 0.12022363, 0.70187986, 0.77832615, 0.82126629, 0.63236392, 0.15563904, 0.09755978, 0.54915464, 0.25058913, 0.45112196],
        [0.45621559, 0.78145081, 0.64627969, 0.56757075, 0.01061873, 0.04715587, 0.28723872, 0.65217435, 0.24890494, 0.61308855, 0.79217201, 0.71212727]])
梯度结果：Tensor(shape=[4, 12], dtype=float32, place=CPUPlace, stop_gradient=False,
       [[0.01919817, 0.01859474, 0.01613311, 0.01882833, 0.01683824, 0.01570455, 0.02077235, 0.01813206, 0.01426661, 0.01618561, 0.02070865, 0.02075072],
        [0.01848637, 0.02083102, 0.02064958, 0.02033491, 0.01930835, 0.01523355, 0.01799664, 0.02080188, 0.01893367, 0.02053295, 0.01664377, 0.02029531],
        [0.01474463, 0.01779640, 0.02068296, 0.01590895, 0.01483520, 0.01419364, 0.01680485, 0.02058151, 0.02073427, 0.01777013, 0.02018264, 0.01874914],
        [0.01870263, 0.01478943, 0.01663187, 0.01756686, 0.02083216, 0.02081018, 0.01997979, 0.01655763, 0.02019131, 0.01703906, 0.01463127, 0.01577028]])
1
2
3
4
5
6
7
8
9
10

为了验证算子的正确性，我们可以跟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, rtol=1e-3, atol=1e-2).numpy())
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

sin_result:  [ True]
sin_grad_result:  [ True]
1
2

从输出结果可以看出，我们自定义的算子从实现功能上来说是正确的。但是还有一些误差，精度并不是特别高。

五、总结与升华

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

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

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

作者简介

北京联合大学 机器人学院 自动化专业 2018级 本科生 郑博培

中国科学院自动化研究所复杂系统管理与控制国家重点实验室实习生

百度飞桨开发者技术专家 PPDE

百度飞桨官方帮帮团、答疑团成员

深圳柴火创客空间 认证会员

百度大脑 智能对话训练师

阿里云人工智能、DevOps助理工程师

我在AI Studio上获得至尊等级，点亮10个徽章，来互关呀！！！

https://aistudio.baidu.com/aistudio/personalcenter/thirdview/147378