自定義CPU算子的實現和使用

• 一、底層原理
• 二、C++自定義算子格式
• • 1.基本格式
  • 2.適配多種數據類型
  • 3.維度與類型的推導
  • 4.自定義算子注冊
• 三、動手實現CPU算子
• • 1.導入必要的頭文件
  • 2.實現forward計算函數
  • 3.實現backward計算函數
  • 4.維度推導
  • 5.自定義算子注冊
• 四、自定義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來完成:

調用relu算子的forward計算函數完成Y的計算

創建backward所需的Op算子以及輸入輸出變量(此時不進行計算,待后續調用backward()后才會進行反向計算)

PaddlePaddle Op算子正反向計算(動態圖模式)

通過上面展示的底層原理，其實不難發現，一個算子最關鍵的部分就是前向傳播與反向計算，這兩個部分是算子的核心。

C++自定義算子內部原理(動態圖模式)

在動態圖模式執行Y=custom_ relu(X)時:使用C++自定義算子與原生算子的執行流程相同。

但和原生算子區別是原生算子隨框架一起編譯；自定義算子可單獨編譯。但最后都會注冊到OpInfoMap中。

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

1.基本格式

在編寫運算函數之前，需要引入 PaddlePaddle 擴展頭文件：

#include "paddle/extension.h"

算子運算函數有特定的函數寫法要求，在編碼過程中需要遵守，基本形式如下：

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

這一部分其實就是固定格式，所有用C++編寫的Paddle算子都需要使用這個格式。換句話說，這是Paddle提供的算子接口，只需要按照這個接口定義算子即可。

2.適配多種數據類型

在實際開發中，一個算子往往需要支持多種數據類型，這時就需要用到模板類。在上面接口上方定義：

template <typename data_t>

需要注意的是：模板參數名 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()), "`"); }

如果不想使用 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));

需要注意的是：

• 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.")

引入 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]);} }

接著只需要將前面實現的計算函數按照前面給的格式套進前向傳播即可：

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}; }

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}; }

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}; }

因為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));

四、自定義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 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)) 前向計算結果：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]])

為了驗證算子的正確性，我們可以跟Paddle現有的算子做對比，看看前向傳播和梯度的計算結果是否一致：

import paddle import paddle.nn.functional as F import numpy as npdevice = '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()) sin_result: [ True] sin_grad_result: [ True]

從輸出結果可以看出，我們自定義的算子從實現功能上來說是正確的。但是還有一些誤差，精度并不是特別高。

五、總結與升華

最后總結一下C++自定義算子最主要的思路，其實就是3點：

forward和backward實現

包裝forward和backward函數并注冊

編譯加載并調用算子

從我的感受來說，我認為第一點是最為重要的部分，特別是反向傳播里梯度的計算，需要一定的數學基礎，要對神經網絡的工作機制有較為深刻的理解。

作者簡介

北京聯合大學 機器人學院 自動化專業 2018級 本科生 鄭博培

中國科學院自動化研究所復雜系統管理與控制國家重點實驗室實習生

百度飛槳開發者技術專家 PPDE

百度飛槳官方幫幫團、答疑團成員

深圳柴火創客空間 認證會員

百度大腦 智能對話訓練師

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

我在AI Studio上獲得至尊等級，點亮10個徽章，來互關呀！！！

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

總結

以上是生活随笔為你收集整理的飞桨高阶使用教程：自定义CPU算子的实现和使用的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。