自定義CPU算子的實(shí)現(xiàn)和使用

• 一、底層原理
• 二、C++自定義算子格式
• • 1.基本格式
  • 2.適配多種數(shù)據(jù)類型
  • 3.維度與類型的推導(dǎo)
  • 4.自定義算子注冊(cè)
• 三、動(dòng)手實(shí)現(xiàn)CPU算子
• • 1.導(dǎo)入必要的頭文件
  • 2.實(shí)現(xiàn)forward計(jì)算函數(shù)
  • 3.實(shí)現(xiàn)backward計(jì)算函數(shù)
  • 4.維度推導(dǎo)
  • 5.自定義算子注冊(cè)
• 四、自定義CPU算子的使用
• 五、總結(jié)與升華
• 作者簡(jiǎn)介

算子（Operator，簡(jiǎn)稱Op）是構(gòu)建神經(jīng)網(wǎng)絡(luò)的基礎(chǔ)組件。在網(wǎng)絡(luò)模型中，算子對(duì)應(yīng)層中的計(jì)算邏輯，例如：卷積層（Convolution Layer）是一個(gè)算子；全連接層（Fully-connected Layer， FC layer）中的權(quán)值求和過(guò)程，是一個(gè)算子。學(xué)會(huì)定制化算子的C++實(shí)現(xiàn)可以更深入地了解神經(jīng)網(wǎng)絡(luò)運(yùn)行的底層邏輯。

張量與算子

• 張量(Tensor)，可以理解為多維數(shù)組,可以具有任意多的維度,不同Tensor可以有不同的數(shù)據(jù)類型(dtype)和形狀(shape)。
• 算子(Operator)也簡(jiǎn)稱為OP，負(fù)責(zé)對(duì)Tensor執(zhí)行各種運(yùn)算處理,可以理解為一個(gè)計(jì)算函數(shù),函數(shù)的輸入和輸出都為Tensor。在PaddlePaddle中定義了大量的Operator來(lái)完成常見神經(jīng)網(wǎng)絡(luò)模型的Tensor運(yùn)算處理,如conv2d, pool2d, Relu等

一、底層原理

PaddlePaddle中所有Op算子都會(huì)注冊(cè)在OpInfoMap中,在Python端調(diào)用Op執(zhí)行運(yùn)算操作時(shí),通過(guò)TraceOp在OpInfoMap找到對(duì)應(yīng)的Op并調(diào)用其計(jì)算kernel ,完成計(jì)算并返回結(jié)果。

換句話說(shuō)，因?yàn)橛布牟煌?#xff0c;相同的算子需要不同的kernel，比如你寫了一個(gè)在CPU上執(zhí)行的算子，那么這個(gè)算子只能在CPU上運(yùn)行，要想在別的計(jì)算平臺(tái)上運(yùn)行，還需要實(shí)現(xiàn)該平臺(tái)的kernel。這篇文章講的主要是怎么實(shí)現(xiàn)一個(gè)能在CPU上運(yùn)行的算子。

PaddlePaddle Op算子體系(動(dòng)態(tài)圖模式)

例如，在動(dòng)態(tài)圖模式執(zhí)行Y=relu(X)時(shí),框架會(huì)通過(guò)TraceOp來(lái)完成:

調(diào)用relu算子的forward計(jì)算函數(shù)完成Y的計(jì)算

創(chuàng)建backward所需的Op算子以及輸入輸出變量(此時(shí)不進(jìn)行計(jì)算,待后續(xù)調(diào)用backward()后才會(huì)進(jìn)行反向計(jì)算)

PaddlePaddle Op算子正反向計(jì)算(動(dòng)態(tài)圖模式)

通過(guò)上面展示的底層原理，其實(shí)不難發(fā)現(xiàn)，一個(gè)算子最關(guān)鍵的部分就是前向傳播與反向計(jì)算，這兩個(gè)部分是算子的核心。

C++自定義算子內(nèi)部原理(動(dòng)態(tài)圖模式)

在動(dòng)態(tài)圖模式執(zhí)行Y=custom_ relu(X)時(shí):使用C++自定義算子與原生算子的執(zhí)行流程相同。

但和原生算子區(qū)別是原生算子隨框架一起編譯；自定義算子可單獨(dú)編譯。但最后都會(huì)注冊(cè)到OpInfoMap中。

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

1.基本格式

在編寫運(yùn)算函數(shù)之前，需要引入 PaddlePaddle 擴(kuò)展頭文件：

#include "paddle/extension.h"

算子運(yùn)算函數(shù)有特定的函數(shù)寫法要求，在編碼過(guò)程中需要遵守，基本形式如下：

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

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

2.適配多種數(shù)據(jù)類型

在實(shí)際開發(fā)中，一個(gè)算子往往需要支持多種數(shù)據(jù)類型，這時(shí)就需要用到模板類。在上面接口上方定義：

template <typename data_t>

需要注意的是：模板參數(shù)名 data_t 用于適配不同的數(shù)據(jù)類型，不可更改為其他命名，否則會(huì)編譯失敗

然后通過(guò) switch-case 語(yǔ)句實(shí)現(xiàn)支持多種數(shù)據(jù)類型的操作：

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 來(lái)實(shí)現(xiàn)，也可以使用官方提供的dispatch宏，如PD_DISPATCH_FLOATING_TYPES

3.維度與類型的推導(dǎo)

PaddlePaddle 框架同時(shí)支持動(dòng)態(tài)圖與靜態(tài)圖的執(zhí)行模式，在靜態(tài)圖模式下，組網(wǎng)階段需要完成 Tensor shape 和 dtype 的推導(dǎo)，從而生成正確的模型描述，用于后續(xù)Graph優(yōu)化與執(zhí)行。因此，除了算子的運(yùn)算函數(shù)之外，還需要實(shí)現(xiàn)前向運(yùn)算的維度和類型的推導(dǎo)函數(shù)。

維度推導(dǎo)（InferShape）和類型推導(dǎo)（InferDtype）的函數(shù)寫法也是有要求的，格式如下：

需要注意的是，輸入輸出參數(shù)與forward計(jì)算函數(shù)的輸入輸出Tensor應(yīng)該按順序一一對(duì)應(yīng)：

對(duì)于僅有一個(gè)輸入Tensor和一個(gè)輸出Tensor的自定義算子，如果輸出Tensor和輸入Tensor的shape和dtype一致,可以省略InferShape和InferDtype函數(shù)的實(shí)現(xiàn)，但其他場(chǎng)景下均需要實(shí)現(xiàn)這兩個(gè)函數(shù)。

4.自定義算子注冊(cè)

最后，需要調(diào)用 PD_BUILD_OP 系列宏，構(gòu)建算子的描述信息，并關(guān)聯(lián)前述算子運(yùn)算函數(shù)和維度、類型推導(dǎo)函數(shù)。其格式如下：

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 用于構(gòu)建前向算子，其括號(hào)內(nèi)為算子名，也是后面在python端使用的接口名，注意前后不需要引號(hào)，注意該算子名不能與 PaddlePaddle 內(nèi)已有算子名重名
• PD_BUILD_GRAD_OP 用于構(gòu)建前向算子對(duì)應(yīng)的反向算子，PD_BUILD_DOUBLE_GRAD_OP 用于構(gòu)建前反向算子對(duì)應(yīng)的二次求導(dǎo)算子。Paddle目前支持的多階導(dǎo)數(shù)只支持到二階導(dǎo)

三、動(dòng)手實(shí)現(xiàn)CPU算子

下面將以一個(gè)比較簡(jiǎn)單的Sin函數(shù)為例，自定義一個(gè)CPU算子。

1.導(dǎo)入必要的頭文件

#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 擴(kuò)展頭文件以及宏定義，檢驗(yàn)輸入的格式。

2.實(shí)現(xiàn)forward計(jì)算函數(shù)

為了適配多種數(shù)據(jù)類型，這里首先加上模板類。

前向計(jì)算最重要的就是實(shí)現(xiàn)計(jì)算函數(shù)，C++里提供了一些基礎(chǔ)運(yùn)算的函數(shù)，可以直接使用，基本語(yǔ)法一般為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]);} }

接著只需要將前面實(shí)現(xiàn)的計(jì)算函數(shù)按照前面給的格式套進(jìn)前向傳播即可：

std::vector<paddle::Tensor> sin_cpu_forward(const paddle::Tensor& x) {// 數(shù)據(jù)準(zhǔn)備CHECK_CPU_INPUT(x);auto out = paddle::Tensor(paddle::PlaceType::kCPU, x.shape()); // 聲明輸出變量out，需傳入兩個(gè)參數(shù)（運(yùn)行的設(shè)備類型及維度信息）// 計(jì)算實(shí)現(xiàn)PD_DISPATCH_FLOATING_TYPES(x.type(), "sin_cpu_forward_kernel", ([&] {sin_cpu_forward_kernel<data_t>( // 調(diào)用前面定義好的前向計(jì)算函數(shù)x.data<data_t>(), // 獲取輸入的內(nèi)存地址，即從內(nèi)存空間中取數(shù)據(jù)out.mutable_data<data_t>(x.place()), x.size()); // 為輸出申請(qǐng)內(nèi)存空間 }));return {out}; }

3.實(shí)現(xiàn)backward計(jì)算函數(shù)

這部分需要一定的數(shù)學(xué)基礎(chǔ)，要了解偏微分的計(jì)算方法，理解神經(jīng)網(wǎng)絡(luò)的梯度概念，我在實(shí)現(xiàn)過(guò)程中也查閱了一些資料，給大家分享：

• 3blue1brown：https://www.3blue1brown.com/lessons/backpropagation-calculus
• 神經(jīng)網(wǎng)絡(luò)之梯度下降法及其實(shí)現(xiàn)
• wolframalpha：https://www.wolframalpha.com/

最后一個(gè)網(wǎng)站是一個(gè)可以直接計(jì)算偏導(dǎo)數(shù)的網(wǎng)站，比較方便，比如這里需要計(jì)算sin函數(shù)的偏導(dǎo):

反向傳播最難的就是計(jì)算梯度，如果會(huì)計(jì)算，其實(shí)就很簡(jiǎn)單了，跟前向計(jì)算是類似的：

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]); // 結(jié)果是返回的梯度值乘函數(shù)導(dǎo)數(shù)值} }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());// 計(jì)算實(shí)現(xiàn)PD_DISPATCH_FLOATING_TYPES(out.type(), "sin_cpu_backward_kernel", ([&] {sin_cpu_backward_kernel<data_t>(grad_out.data<data_t>(), // 獲取內(nèi)存地址，即從內(nèi)存空間中取數(shù)據(jù)out.data<data_t>(), // 獲取內(nèi)存地址，即從內(nèi)存空間中取數(shù)據(jù)grad_x.mutable_data<data_t>(x.place()), // 申請(qǐng)內(nèi)存空間out.size()); // 傳入輸出的維度信息}));return {grad_x}; }

4.維度推導(dǎo)

維度推導(dǎo)部分其實(shí)只需要根據(jù)格式實(shí)現(xiàn)InferShape和InferDtype函數(shù)即可：

// 維度推導(dǎo) std::vector<std::vector<int64_t>> sinInferShape(std::vector<int64_t> x_shape) {return {x_shape}; }// 類型推導(dǎo) std::vector<paddle::DataType> sinInferDtype(paddle::DataType x_dtype) {return {x_dtype}; }

因?yàn)閟in(x)函數(shù)輸入和輸出的維度一致，所以可以省略InferShape和InferDtype函數(shù)的實(shí)現(xiàn)。

5.自定義算子注冊(cè)

最后也是按照格式完成自定義算子的注冊(cè)即可：

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 (即時(shí)編譯)安裝加載自定義算子，其基本格式如下：

在本項(xiàng)目中，我已經(jīng)將算子寫好，位于custom_op/custom_sin_cpu.cc，直接調(diào)用即可：

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.....

使用該算子也非常簡(jiǎn)單，直接使用即可，如下所示：

import paddle import paddle.nn.functional as F import numpy as np# 定義執(zhí)行環(huán)境 device = 'cpu' paddle.set_device(device)# 將輸入數(shù)據(jù)轉(zhuǎn)換為張量 data = np.random.random([4, 12]).astype(np.float32) x = paddle.to_tensor(data, stop_gradient=False)# 調(diào)用自定義算子實(shí)現(xiàn)前向計(jì)算 y = custom_sin_cpu(x) # 調(diào)用自定義算子實(shí)現(xiàn)反向傳播 y.mean().backward()print("前向計(jì)算結(jié)果：{}".format(y)) print("梯度結(jié)果：{}".format(x.grad)) 前向計(jì)算結(jié)果：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]]) 梯度結(jié)果：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]])

為了驗(yàn)證算子的正確性，我們可以跟Paddle現(xiàn)有的算子做對(duì)比，看看前向傳播和梯度的計(jì)算結(jié)果是否一致：

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表示結(jié)果正確 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]

從輸出結(jié)果可以看出，我們自定義的算子從實(shí)現(xiàn)功能上來(lái)說(shuō)是正確的。但是還有一些誤差，精度并不是特別高。

五、總結(jié)與升華

最后總結(jié)一下C++自定義算子最主要的思路，其實(shí)就是3點(diǎn)：

forward和backward實(shí)現(xiàn)

包裝forward和backward函數(shù)并注冊(cè)

編譯加載并調(diào)用算子

從我的感受來(lái)說(shuō)，我認(rèn)為第一點(diǎn)是最為重要的部分，特別是反向傳播里梯度的計(jì)算，需要一定的數(shù)學(xué)基礎(chǔ)，要對(duì)神經(jīng)網(wǎng)絡(luò)的工作機(jī)制有較為深刻的理解。

作者簡(jiǎn)介

北京聯(lián)合大學(xué) 機(jī)器人學(xué)院 自動(dòng)化專業(yè) 2018級(jí) 本科生 鄭博培

中國(guó)科學(xué)院自動(dòng)化研究所復(fù)雜系統(tǒng)管理與控制國(guó)家重點(diǎn)實(shí)驗(yàn)室實(shí)習(xí)生

百度飛槳開發(fā)者技術(shù)專家 PPDE

百度飛槳官方幫幫團(tuán)、答疑團(tuán)成員

深圳柴火創(chuàng)客空間 認(rèn)證會(huì)員

百度大腦 智能對(duì)話訓(xùn)練師

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

我在AI Studio上獲得至尊等級(jí)，點(diǎn)亮10個(gè)徽章，來(lái)互關(guān)呀！！！

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

總結(jié)

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

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。