LayerNorm Plugin的使用與說明

- 前言
- 0. 簡述
- 1. Layernorm Plugin的使用
- - 1.1 源碼下載
  - 1.2 模型下載和修改
  - 1.3 環境配置
  - 1.4 編譯
  - 1.4 engine生成和執行(trtexec)
  - 1.5 enging生成和執行(C++ API)
- 2. 補充說明
- - 2.1 RTMO顯存占用問題
  - 2.2 插件找不到的說明
  - 2.3 LayerNorm plugin封裝的嘗試
  - 2.4 layerNorm plugin核函數實現淺析
- 結語
- 下載鏈接
- 參考

前言

最近在 CUDA-BEVFusion 項目上看到杜老師有添加 layernorm plugin 的支持，這里分享博主在使用 layernorm plugin 時做的一些嘗試，并不涉及任何原理性的分析，若有問題歡迎各位看官批評指正😄

0. 簡述

前面我們在做模型部署的時候經常會遇到 LayerNormalization 這個算子，比如 RTMO、RT-DETR 等模型都包含該算子，而 tensorRT 官方只在 8.6 版本才提供了該算子的支持，那如果低版本的 tensorRT 我們依舊想解析該算子是不是可以考慮寫插件來支持呢

因此這里我們就來做這么一件事情，把 layernorm plugin 添加到我們的項目中并使用它完成模型推理，我們以 tensorRT_Pro-YOLOv8 這個 repo 為例，來添加 layernorm plugin 的支持

整個實現流程如下：

tensorRT_Pro-YOLOv8 項目下載，添加 custom_layernorm.cu 插件
ONNX 模型下載（RTMO 或 RT-DETR）
修改 ONNX 模型 LayerNorm 節點并保存（主要修改其 op_type）
將 layernorm plugin 編譯成動態庫 libcustom_layernorm.so
加載 libcustom_layernorm.so 利用 C++ API 或 trtexec 工具生成 engine
利用生成的 engine 完成推理

下面我們就來按照這個流程走一遍

1. Layernorm Plugin的使用

1.1 源碼下載

tensorRT_Pro-YOLOv8 的代碼可以直接從 GitHub 官網上下載，源碼下載地址是 https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8，Linux 下代碼克隆指令如下：

git clone https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8.git

也可手動點擊下載，點擊右上角的 Code 按鍵，將代碼下載下來。至此整個項目就已經準備好了。

之后我們需要添加 layernorm plugin 的實現，在 tensorRT_Pro-YOLOv8/src/tensorRT/onnxplugin/plugins 文件夾下新建 custom_layernorm.cu 文件，其內容如下：

#include <NvInfer.h>
#include <NvInferPlugin.h>
#include <vector>
#include <string>
#include <assert.h>
#include <cuda_fp16.h>using namespace nvinfer1;template<typename T>
static void __global__ layernorm_kernel(const T* x, const T* weight, const T* bias, T* y, int N, int C, float epsilon);template<>
void __global__ layernorm_kernel<float>(const float* x, const float* weight, const float* bias, float* y, int N, int C, float epsilon){int idx = blockIdx.y * blockDim.y + threadIdx.y;if(idx >= N) return;// x: N, C// y: N, C// weight: C// bias:   Cconst float* px = x + idx * C;float*       py = y + idx * C;// reduce sumfloat sq = 0.0f;float s  = 0.0f;float diver = 1.0f / C;for(int ic = threadIdx.x; ic < C; ic += warpSize){float x = px[ic];s += x;sq = fmaf(x, x * diver, sq);}for (int mask = 16; mask > 0; mask /= 2)s += __shfl_xor_sync(0xffffffff, s, mask);for (int mask = 16; mask > 0; mask /= 2)sq += __shfl_xor_sync(0xffffffff, sq, mask);float mean = s / C;float rstd = rsqrtf(sq - mean * mean + epsilon);for(int ic = threadIdx.x; ic < C; ic += warpSize) py[ic] = (px[ic] - mean) * weight[ic] * rstd + bias[ic];
}template<>
void __global__ layernorm_kernel<half>(const half* x, const half* weight, const half* bias, half* y, int N, int C, float epsilon){int idx = blockIdx.y * blockDim.y + threadIdx.y;if(idx >= N) return;// x: N, C// y: N, C// weight: C// bias:   Cconst half* px = x + idx * C;half* py = y + idx * C;// reduce sumfloat sq = 0.0f;float s  = 0.0f;float diver = 1.0f / C;for(int ic = threadIdx.x; ic < C; ic += warpSize){float x = __half2float(px[ic]);s += x;sq = fmaf(x, x * diver, sq);}for (int mask = 16; mask > 0; mask /= 2)s += __shfl_xor_sync(0xffffffff, s, mask);for (int mask = 16; mask > 0; mask /= 2)sq += __shfl_xor_sync(0xffffffff, sq, mask);float mean = s / C;float rstd = rsqrtf(sq - mean * mean + epsilon);// for(int ic = threadIdx.x; ic < C; ic += warpSize) // py[ic] = __float2half((__half2float(px[ic]) - mean) * __half2float(weight[ic]) * rstd) + bias[ic];// ===== modify =====for(int ic = threadIdx.x; ic < C; ic += warpSize) {float px_val = __half2float(px[ic]);float weight_val = __half2float(weight[ic]);float bias_val = __half2float(bias[ic]);float result = (px_val - mean) * weight_val * rstd + bias_val;py[ic] = __float2half(result);}        
}class LayerNormPlugin : public IPluginV2DynamicExt{
public:float epsilon;int axis;// construct by creatationLayerNormPlugin(float epsilon, int axis){this->epsilon = epsilon;this->axis    = axis;}// construct by deserializationLayerNormPlugin(const void* data, size_t size){const unsigned char* pdata = (const unsigned char*)data;this->epsilon = *(float*)pdata;  pdata += sizeof(this->epsilon);this->axis    = *((int*)pdata);}IPluginV2DynamicExt* clone() const noexcept override{return new LayerNormPlugin(this->epsilon, this->axis);}virtual DimsExprs getOutputDimensions(int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept{return inputs[0];}virtual bool supportsFormatCombination(int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept{return inOut[pos].format == TensorFormat::kLINEAR && (inOut[pos].type == DataType::kFLOAT || inOut[pos].type == DataType::kHALF) && inOut[pos].type == inOut[0].type;}virtual void configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs,DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept{}virtual size_t getWorkspaceSize(PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs,int32_t nbOutputs) const noexcept{return 0;}virtual int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept{// B, N, Cint N = inputDesc[0].dims.d[0] * inputDesc[0].dims.d[1];int C = inputDesc[0].dims.d[2];const void* x      = inputs[0];const void* weight = inputs[1];const void* bias   = inputs[2];void* y            = outputs[0];dim3 block(32, 8);dim3 grid(1, (N + block.y - 1) / block.y);if(inputDesc[0].type == DataType::kHALF){layernorm_kernel<half><<<grid, block, 0, stream>>>((half*)x, (half*)weight, (half*)bias, (half*)y, N, C, this->epsilon);}else if(inputDesc[0].type == DataType::kFLOAT){layernorm_kernel<float><<<grid, block, 0, stream>>>((float*)x, (float*)weight, (float*)bias, (float*)y, N, C, this->epsilon);}else{// not implementedreturn 1;}return 0;}virtual nvinfer1::DataType getOutputDataType(int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept{return inputTypes[0];}virtual int32_t initialize() noexcept{return 0;}virtual void terminate() noexcept{}virtual void serialize(void* buffer) const noexcept{unsigned char* pdata = (unsigned char*)buffer;*(float*)pdata = this->epsilon;  pdata += sizeof(this->epsilon);*(int*)pdata   = this->axis;}virtual void destroy() noexcept{}virtual void setPluginNamespace(AsciiChar const* pluginNamespace) noexcept{}virtual AsciiChar const* getPluginNamespace() const noexcept{return "";}virtual AsciiChar const* getPluginType() const noexcept{return "CustomLayerNormalization";}virtual AsciiChar const* getPluginVersion() const noexcept{return "1";}virtual int32_t getNbOutputs() const noexcept {return 1;}virtual size_t getSerializationSize() const noexcept{return sizeof(this->epsilon) + sizeof(this->axis);}
};class LayerNormCreater : public IPluginCreator{
public:std::vector<PluginField> fields;PluginFieldCollection field_collection;std::string namespace_name = "ours";LayerNormCreater(){fields = {PluginField{"epsilon", nullptr, PluginFieldType::kFLOAT32, 1},PluginField{"axis",    nullptr, PluginFieldType::kINT32, 1},};field_collection.fields   = fields.data();field_collection.nbFields = fields.size();}virtual AsciiChar const* getPluginName() const noexcept{return "CustomLayerNormalization";}virtual AsciiChar const* getPluginVersion() const noexcept{return "1";}virtual PluginFieldCollection const* getFieldNames() noexcept{return &field_collection;}virtual IPluginV2* createPlugin(AsciiChar const* name, PluginFieldCollection const* fc) noexcept{assert(strcmp("epsilon", fc->fields[0].name) == 0);assert(strcmp("axis",    fc->fields[1].name) == 0);float epsilon = *(float*)(fc->fields[0].data);int axis      = *(int*)(fc->fields[1].data);printf("epsilon = %g, axis=%d\n", epsilon, axis);return new LayerNormPlugin(epsilon, axis);}virtual IPluginV2* deserializePlugin(AsciiChar const* name, void const* serialData, size_t serialLength) noexcept{return new LayerNormPlugin(serialData, serialLength);}virtual void setPluginNamespace(AsciiChar const* pluginNamespace) noexcept{}virtual AsciiChar const* getPluginNamespace() const noexcept{return "";}
};REGISTER_TENSORRT_PLUGIN(LayerNormCreater);

代碼完全 copy 自杜老師的實現，大家可以查看：CUDA-BEVFusion/src/plugins/custom_layernorm.cu

這里有一點需要注意，博主對代碼進行了一些簡單修改，這是因為博主在后續的編譯過程中出現了如下問題：

在這里插入圖片描述

錯誤信息顯示關于 __half 類型到基本類型的多個可能轉換，這個問題出現在 custom_layernorm.cu 文件的第 78 行。當 CUDA 編譯器遇到一個 __half 類型的值，并嘗試將它轉換到一個內建類型時，它發現有多個轉換函數可用，導致編譯器無法決定使用哪一個，因此報錯

不過奇怪的是代碼中的變量明確使用了 __half2float，按理來說不會出現這個問題，博主也不知道問題出現在哪，按照 ChatGPT 的提示進行了修改，先將所有的變換轉換為 float 進行計算后再轉回 half，具體修改位置是在 layernorm fp16 核函數的實現代碼中：

// for(int ic = threadIdx.x; ic < C; ic += warpSize) 
//     py[ic] = __float2half((__half2float(px[ic]) - mean) * __half2float(weight[ic]) * rstd) + bias[ic];
// ===== modify =====
for(int ic = threadIdx.x; ic < C; ic += warpSize) {float px_val = __half2float(px[ic]);float weight_val = __half2float(weight[ic]);float bias_val = __half2float(bias[ic]);float result = (px_val - mean) * weight_val * rstd + bias_val;py[ic] = __float2half(result);
}

這樣做似乎有點呆，因為 bias_val 變量本身就是 half 類型的，先轉成 float 類型計算完后又再轉回了 half，不過這樣確實可以解決編譯問題

1.2 模型下載和修改

我們測試的模型可以選擇 RTMO 或者 RT-DETR，因為這兩個模型都包含 LayerNorm 節點，博主這里準備了導出好的 ONNX 模型，大家可以點擊 here 進行下載，將下載好的模型放在 tensorRT_Pro-YOLOv8/workspace 文件夾下，方便后續 engine 的生成

關于 RT-DETR 和 RTMO 的 ONNX 導出大家可以參考 RT-DETR推理詳解及部署實現和 MMPose-RTMO推理詳解及部署實現（上），這里博主就不再贅述了。

在拿到 ONNX 之后我們需要做一些修改，因為原始導出的 LayerNorm 的 op_type 是 LayerNormalization，我們需要修改為我們自定義的 op_type，從 custom_layernorm.cu 代碼中我們不難看出自定義的 LayerNorm 的 op_type 為 CustomLayerNormalization，因此我們將原始的 LayerNormalization 修改為 CustomLayerNormalization 即可，這個我們在杜老師的課程中有簡單講過，大家感興趣的可以看看：5.4.tensorRT基礎(2)-學習第一個插件的編寫

這里我們可以使用 onnx_graphsurgeon 這個庫來完成這個操作，它是 NVIDIA 提供的一個創建和修改 ONNX 的工具，我們可以使用如下指令進行安裝：

python3 -m pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com

關于 onnx-graph-surgeon 的使用我們在韓君老師的課程中有簡單講過，大家感興趣的可以看看：三. TensorRT基礎入門-onnx-graph-surgeon

我們在 tensorRT_Pro-YOLOv8/workspace 文件夾下新建一個 change_layernorm.py 文件，其內容如下：(from ChatGPT)

import onnx
import onnx_graphsurgeon as gs# 加載 ONNX 模型
input_model_path = "rtmo-s_8xb32-600e_body7-640x640.onnx"
output_model_path = "rtmo-s_8xb32-600e_body7-640x640.plugin.onnx"
graph = gs.import_onnx(onnx.load(input_model_path))# 遍歷圖中的所有節點
for node in graph.nodes:if node.op == "LayerNormalization":node.op = "CustomLayerNormalization"# 添加自定義屬性node.attrs["name"] = "LayerNormPlugin"node.attrs["info"] = "This is custom LayerNormalization node"# 刪除無用的節點和張量
graph.cleanup()# 導出修改后的模型
onnx.save(gs.export_onnx(graph), output_model_path)

執行該代碼后會在當前目錄生成 rtmo-s_8xb32-600e_body7-640x640.plugin.onnx，一個修改過的 ONNX，這個 ONNX 將被用于我們后續的 engine 構建

在上述代碼中我們遍歷了整個 Graph 的節點并找到所有的 LayerNormalization，將其修改為 CustomLayerNormalization，此外我們還添加了一些自定義的屬性，當然這些屬性我們并沒有使用到，在后面插件的封裝中我們會使用到。這里大家要有一個概念，那就是生成好的 ONNX 模型是可編輯，可修改的，我們可以把它當成一個記事本，這也是杜老師反復提到的概念

那博主這邊還有一個顧慮，那就是我們是不是不應該簡單的使用 onnx_graphsurgeon 修改 op_type 呢，我們是不是應該在 .pt 模型導出 ONNX 時利用 symbolic 函數將標準的 LayerNorm 層替換為一個自定義的 LayerNorm 實現呢🤔

在 CUDA-BEVFusion 中的 CUDA-BEVFusion/qat/export-transfuser.py#L242 就是這么做的，具體代碼如下所示：

class CustomLayerNormImpl(torch.autograd.Function):@staticmethoddef forward(ctx, input, normalized_shape, weight, bias, eps, x_shape):return F.layer_norm(input, normalized_shape, weight, bias, eps)@staticmethoddef symbolic(g, input, normalized_shape, weight, bias, eps, x_shape):y = g.op("nv::CustomLayerNormalization", input, weight, bias, axis_i=-1, epsilon_f=eps)y.setType(input.type().with_sizes(x_shape))return yclass CustomLayerNorm(nn.LayerNorm):def forward(self, input: torch.Tensor) -> torch.Tensor:return CustomLayerNormImpl.apply(input, self.normalized_shape, self.weight, self.bias, self.eps, input.size())@staticmethoddef convert(old: nn.LayerNorm):Custom_layernorm = CustomLayerNorm(old.normalized_shape, old.eps, old.elementwise_affine)if Custom_layernorm.weight is not None:Custom_layernorm.weight.data = old.weight.dataCustom_layernorm.bias.data   = old.bias.datareturn Custom_layernormdef replace_layernorm(model):for name, module in model.named_modules():if isinstance(module, nn.LayerNorm):parent, child = name.rsplit(".", 1)parent = model.get_submodule(parent)setattr(parent, child, CustomLayerNorm.convert(module))

那當然后續博主在測試的過程中也能正常加載插件完成推理，因此也就沒有去做 .pt 導出 ONNX 時替換 LayerNorm 這件事情了

我們下面來看看我們修改后的 ONNX 模型的 LayerNorm 節點發生了什么變化呢？

我們先看原始的 LayerNorm 節點，如下圖所示：

在這里插入圖片描述

從圖中我們可以看到 LayerNorm 這個節點的 op_type 為 LayerNormalization，它的屬性包含 epsilon 和 axis，它的輸入包含 input、weight 和 bias 三部分，輸出只有 output 一個部分

我們再來看看修改后的 LayerNorm 節點，如下圖所示：

在這里插入圖片描述

從圖中我們可以看到原來的 op_type 從 LayerNormalization 修改為 CustomLayerNormalization 了，另外新增了 name 和 info 兩個屬性，這都是我們在代碼中實現的，其它的部分倒是沒什么變化

OK，至此我們完成了項目和模型的下載以及 ONNX 模型的修改，下面我們來看看環境配置

1.3 環境配置

tensorRT_Pro_YOLOv8 這個 repo 的運行需要使用的軟件環境有 TensorRT、CUDA、cuDNN、OpenCV、Protobuf，所有軟件環境的安裝可以參考 Ubuntu20.04軟件安裝大全，這里不再贅述，需要各位看官自行配置好相關環境😄，外網訪問較慢，這里提供下博主安裝過程中的軟件安裝包下載鏈接 Baidu Drive【pwd:yolo】🚀🚀🚀

tensorRT_Pro-YOLOv8 提供 CMakeLists.txt 和 Makefile 兩種方式編譯，我們這里統一使用 CMakeLists.txt

主要修改六處

1. 修改第 13 行，修改 OpenCV 路徑

set(OpenCV_DIR   "/usr/local/include/opencv4/")

2. 修改第 15 行，修改 CUDA 路徑

set(CUDA_TOOLKIT_ROOT_DIR     "/usr/local/cuda-11.6")

3. 修改第 16 行，修改 cuDNN 路徑

set(CUDNN_DIR    "/usr/local/cudnn8.4.0.27-cuda11.6")

4. 修改第 17 行，修改 tensorRT 路徑（版本小于 8.6）

set(TENSORRT_DIR "/home/jarvis/lean/TensorRT-8.5.1.7")

5. 修改第 20 行，修改 protobuf 路徑

set(PROTOBUF_DIR "/home/jarvis/protobuf")

6. 修改第 60 行，新增 layernorm plugin 動態庫編譯內容

########################## custom_layernorm.so ################################
cuda_add_library(custom_layernorm SHAREDsrc/tensorRT/onnxplugin/plugins/custom_layernorm.cu
)target_link_libraries(custom_layernormlibnvinfer.solibnvinfer_plugin.so
)

Note：這里博主使用的 tensorRT 版本為 8.5.1.7，前面提到過 tensorRT 在 8.6 版本之后就已經支持 LayerNormalization 算子了，如下圖所示，我們的目的就是在低版本的 tensorRT 中實現該節點的推理與解析，因此我們這里拿低版本的 tensorRT 進行驗證測試

在這里插入圖片描述

1.4 編譯

將上述環境配置好后我們就來編譯整個項目，目的是生成 libcustom_layernorm.so 動態庫方便后續加載并生成 engine

整個編譯流程如下：

cd tensorRT_Pro-YOLOv8
mkdir build && cd build
cmake .. && make -j64
cp libcustom_layernorm.so ../workspace

部分輸出如下圖所示：

在這里插入圖片描述

編譯完成后我們可以在 build 目錄下看到 libcustom_layernorm.so 動態庫，并將這個動態庫復制到了 tensorRT_Pro-YOLOv8/workspace 目錄下

1.4 engine生成和執行(trtexec)

engine 的生成我們可以使用 trtexec 工具生成或者使用 C++ API 接口生成，這里我們先看 trtexec 工具生成 engine 的方式

在開始之前我們可以先測試下 TensorRT-8.5.1.7 的 trtexec 能否成功構建 engine，在 tensorRT_Pro-YOLOv8/workspace 文件夾下新建腳本文件 build.sh，其內容如下：

#! /usr/bin/bashTRTEXEC=/home/jarvis/lean/TensorRT-8.5.1.7/bin/trtexec# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/jarvis/lean/TensorRT-8.5.1.7/lib${TRTEXEC} \--onnx=rtmo-s_8xb32-600e_body7-640x640.onnx \--minShapes=images:1x3x640x640 \--optShapes=images:1x3x640x640 \--maxShapes=images:4x3x640x640 \--memPoolSize=workspace:2048 \--saveEngine=rtmo-s_8xb32-600e_body7-640x640.FP32.trtmodel

在終端通過 bash build.sh 可以運行該腳本，大家會看到如下的日志信息：

在這里插入圖片描述

從日志信息中我們可以看到 tensorRT(8.5.1.7) 不支持 LayerNormalization 這個 op，于是它嘗試將 LayerNormalization 作為插件導入并開始搜索插件庫，但是它發現插件庫中也沒有該算子的實現，因此最終報錯提示無法解析該 ONNX 文件

接著我們就來加載插件庫，并解析我們自定義的 LayerNorm 算子生成 engine，在 tensorRT_Pro-YOLOv8/workspace 文件夾下修改腳本文件 build.sh，其內容如下：

#! /usr/bin/bashTRTEXEC=/home/jarvis/lean/TensorRT-8.5.1.7/bin/trtexec# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/jarvis/lean/TensorRT-8.5.1.7/lib${TRTEXEC} \--onnx=rtmo-s_8xb32-600e_body7-640x640.plugin.onnx \--plugins=libcustom_layernorm.so \--minShapes=images:1x3x640x640 \--optShapes=images:1x3x640x640 \--maxShapes=images:4x3x640x640 \--memPoolSize=workspace:2048 \--saveEngine=rtmo-s_8xb32-600e_body7-640x640.plugin.FP32.trtmodel \> trtexec_output.log 2>&1

我們可以看到這里新增了 –plugins 參數，同時將我們之前修改過的 ONNX 放了進來，這里有一些參數的使用我們簡單說明下：

–plugins 必須指定，不然在 tensorRT 構建 engine 的時候會無法解析 LayerNorm 節點
–minShapes、–optShapes 以及 –maxShapes 是動態 shape 模型需要指定的，盡量將 maxBatch 設置得稍微小點不然可能會出現 out of memory 顯存不足的問題，后續我們會來分析
–memPoolSize 用于設置構建 engine 時的內存池大小，這個最好指定可以避免顯存溢出

此外注意將 TRTEXEC 的路徑修改為你自己的路徑，接著在終端執行如下指令即可：

bash build.sh

我們可以在終端看到一些日志信息，如下圖所示：

在這里插入圖片描述

我們從日志信息可以看到模型有正常加載 layernorm plugin 插件并利用它完成了推理，在創建 Plugin 時我們還可以看到 epsilon 和 axis 的打印輸出，這是因為我們在 createPlugin 函數中加入了這兩個參數的打印輸出

等待一段時間后即可完成模型的編譯工作，在 workspace 目錄下會生成 rtmo-s_8xb32-600e_body7-640x640.plugin.FP32.trtmodel 這個 engine 文件

engine 模型生成后我們就可以開始執行了，在執行之前我們需要修改下 src/application/app_rtmo.cpp 中的模型名字，同時將 test_batch_size 設置小點，如下所示：

int test_batch_size = 4; // 176行, 將 batch size 修改為 4test(TRT::Mode::FP32, "rtmo-s_8xb32-600e_body7-640x640.plugin")	// 292行, 將模型名修改為"rtmo-s_8xb32-600e_body7-640x640.plugin"

接著我們就要來開始執行，看下輸出結果，指令如下：

cd tensorRT_Pro-YOLOv8/build
make -j64
cd ../workspace
./pro rtmo

輸出如下圖所示：

在這里插入圖片描述

此外在 workspace 目錄下會生成 rtmo-s_8xb32-600e_body7-640x640.plugin_RTMO_FP32_result 文件夾，該文件夾下保存著推理的圖片，如下所示：

在這里插入圖片描述

可以看到模型正常推理了，這說明我們正常加載了插件并實現了推理

1.5 enging生成和執行(C++ API)

接著我們再來看看利用 C++ API 來加載插件生成 engine，開始之前我們依舊需要簡單修改下 src/application/app_rtmo.cpp 文件的內容，如下所示：

int test_batch_size = 4; // 176行, 將 batch size 修改為 4test(TRT::Mode::FP32, "rtmo-s_8xb32-600e_body7-640x640.plugin")	// 292行, 將模型名修改為"rtmo-s_8xb32-600e_body7-640x640.plugin"

接著我們開始運行生成 engine，指令如下：

cd tensorRT_Pro-YOLOv8/build
make -j64
cd ../workspace
./pro rtmo

輸出如下所示：

在這里插入圖片描述

可以看到我們利用 C++ API 構建 engine，同時看到了 epsilon 和 axis 打印信息，這說明 tensorRT 在創建 plugin，輸出了兩行打印信息也說明 RTMO 這個 model 里面有兩個 LayerNorm 的節點

推理結果如下圖所示：

在這里插入圖片描述

2. 補充說明

2.1 RTMO顯存占用問題

關于 RTMO 構建 engine 時顯存占用問題之前就有看官提過，具體可以看：tensorRT_Pro-YOLOv8/issues/27

這個模型看著不大，但是在構建 engine 時占用顯存非常大，針對原始的 rtmo 模型我們設置下不同的 maxBatch 看下顯存占用情況：

在這里插入圖片描述

batch=1

在這里插入圖片描述

batch=4

在這里插入圖片描述

batch=16

在這里插入圖片描述

batch=16(fp16)

從上圖中我們可以發現 RTMO 這個模型在構建 engine 時占用的顯存確實多，batch=1 的情況下就占用了 2G 的顯存，在 batch=16 的情況下高達 10G 顯存，而且博主測試 batch=16 時 engine 執行會失敗，如下圖所示：

在這里插入圖片描述

可以看到 tensorRT 需要 14G 的顯存去做一些策略優化，但是內存分配失敗，接著在創建 context 執行期間顯示 OutOfMemory 而報錯

那為什么顯存占用會這么高呢，博主這邊也沒有一個確定的結論，不過這個可能和 RTMO 模型的架構有關，里面有一些 transformer 的結構，可能需要提前存儲一些 tensor 之類的變量，也可能和 tensorRT 的優化策略相關，tensorRT 在做 kernel auto-tuning 的時候針對 LayerNorm 等節點可能需要分配大量的內存來做優化

2.2 插件找不到的說明

大家如果遇到了插件找不到的問題，可以嘗試下使用 dlopen(xxx, RTLD_NOW) 的方式手動加載，示例使用如下：

// src/main.cpp
#include <dlfcn.h>int main(int argc, char** argv){const char* method = "yolo";if(argc > 1){method = argv[1];}dlopen("libcustom_layernorm.so", RTLD_NOW);...
}

在 CUDA-BEVFusion 中也是這么做的，具體大家可以參考：CUDA-BEVFusion/src/main.cpp#L231

此外大家如果使用 dlopen 還需要鏈接 libdl.so 庫，在 CMakeLists.txt 新增如下內容即可：

target_link_libraries(pro dl)

2.3 LayerNorm plugin封裝的嘗試

那其實博主最開始嘗試使用 LayerNorm plugin 并不是上面的未修改的代碼，因為杜老師在 tensorRT_Pro 中已經幫我們將插件封裝好了，因此我們只需要按照提供的示例寫 enqueue 函數就行

博主嘗試封裝的代碼如下所示：

#include <onnxplugin/onnxplugin.hpp>
#include <cuda_fp16.h>using namespace ONNXPlugin;template<typename T>
static void __global__ layernorm_kernel(const T* x, const T* weight, const T* bias, T* y, int N, int C, float epsilon);template<>
void __global__ layernorm_kernel<float>(const float* x, const float* weight, const float* bias, float* y, int N, int C, float epsilon){int idx = blockIdx.y * blockDim.y + threadIdx.y;if(idx >= N) return;// x: N, C// y: N, C// weight: C// bias:   Cconst float* px = x + idx * C;float*       py = y + idx * C;// reduce sumfloat sq = 0.0f;float s  = 0.0f;float diver = 1.0f / C;for(int ic = threadIdx.x; ic < C; ic += warpSize){float x = px[ic];s += x;sq = fmaf(x, x * diver, sq);}for (int mask = 16; mask > 0; mask /= 2)s += __shfl_xor_sync(0xffffffff, s, mask);for (int mask = 16; mask > 0; mask /= 2)sq += __shfl_xor_sync(0xffffffff, sq, mask);float mean = s / C;float rstd = rsqrtf(sq - mean * mean + epsilon);for(int ic = threadIdx.x; ic < C; ic += warpSize) py[ic] = (px[ic] - mean) * weight[ic] * rstd + bias[ic];
}template<>
void __global__ layernorm_kernel<half>(const half* x, const half* weight, const half* bias, half* y, int N, int C, float epsilon){int idx = blockIdx.y * blockDim.y + threadIdx.y;if(idx >= N) return;// x: N, C// y: N, C// weight: C// bias:   Cconst half* px = x + idx * C;half* py = y + idx * C;// reduce sumfloat sq = 0.0f;float s  = 0.0f;float diver = 1.0f / C;for(int ic = threadIdx.x; ic < C; ic += warpSize){float x = __half2float(px[ic]);s += x;sq = fmaf(x, x * diver, sq);}for (int mask = 16; mask > 0; mask /= 2)s += __shfl_xor_sync(0xffffffff, s, mask);for (int mask = 16; mask > 0; mask /= 2)sq += __shfl_xor_sync(0xffffffff, sq, mask);float mean = s / C;float rstd = rsqrtf(sq - mean * mean + epsilon);// for(int ic = threadIdx.x; ic < C; ic += warpSize) //     py[ic] = __float2half((__half2float(px[ic]) - mean) * __half2float(weight[ic]) * rstd) + bias[ic];// ===== modify =====for(int ic = threadIdx.x; ic < C; ic += warpSize) {float px_val = __half2float(px[ic]);float weight_val = __half2float(weight[ic]);float bias_val = __half2float(bias[ic]);float result = (px_val - mean) * weight_val * rstd + bias_val;py[ic] = __float2half(result);}
}class LayerNormPlugin : public TRTPlugin{
public:SetupPlugin(LayerNormPlugin);virtual void config_finish() override{}virtual std::shared_ptr<LayerConfig> new_config() override{auto cfg = TRTPlugin::new_config();cfg->support_dtype_set_ = {nvinfer1::DataType::kHALF, nvinfer1::DataType::kFLOAT};return cfg;}virtual nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex, const nvinfer1::DimsExprs* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override{return inputs[0];}int enqueue(const std::vector<GTensor>& inputs, std::vector<GTensor>& outputs, const std::vector<GTensor>& weights, void* workspace, cudaStream_t stream) override{// bx400x256float epsilon = 1e-5;// B, N, Cint N = inputs[0].shape_[0] * inputs[0].shape_[1];int C = inputs[0].shape_[2];const void* x      = inputs[0].ptr_;const void* weight = inputs[1].ptr_;const void* bias   = inputs[2].ptr_;void* y            = outputs[0].ptr_;dim3 block(32, 8);dim3 grid(1, (N + block.y - 1) / block.y);if(config_->usage_dtype_ == TRT::DataType::Float){// fp32layernorm_kernel<float><<<grid, block, 0, stream>>>((float*)x, (float*)weight, (float*)bias, (float*)y, N, C, epsilon);}else if(config_->usage_dtype_ == TRT::DataType::Float16){// fp16layernorm_kernel<half><<<grid, block, 0, stream>>>((half*)x, (half*)weight, (half*)bias, (half*)y, N, C, epsilon);}else{// not implementedreturn 1;}return 0;}
};RegisterPlugin(LayerNormPlugin);

博主這里主要參考 custom_layernorm.cu 和示例 HSwish.cu，可以看到我們創建了一個 LayerNormPlugin 的類繼承自 TRTPlugin，然后在這里面實現一些方法就行，主要就是實現 enqueue 函數，在推理時根據不同的 DataType 調用不同的核函數進行計算，最后調用 RegisterPlugin 將插件注冊即可，實現上有一點不同的是核函數需要的 epsilon 參數博主是直接給定的，并沒有像 custom_layernorm.cu 一樣通過構造函數傳入

插件代碼封裝好后，我們還需要修改下 ONNX 模型，跟之前修改的地方不同，這里主要著重修改兩個點：

LayerNorm 的 op_type 必須修改為 Plugin
LayerNorm 必須新增一個 name 的屬性，類型為 string，名字為 RegisterPlugin 注冊的插件名，這里是 LayerNormPlugin

這樣修改的原因主要是我們對插件進行了封裝，在 src/tensorRT/onnx_parser/builtin_op_importers.cpp 代碼中我們可以看到如下代碼：

DEFINE_BUILTIN_OP_IMPORTER(Plugin)
{std::vector<nvinfer1::ITensor*> inputTensors;std::vector<onnx2trt::ShapedWeights> weights;for(int i = 0; i < inputs.size(); ++i){auto& item = inputs.at(i);if(item.is_tensor()){nvinfer1::ITensor* input = &convertToTensor(item, ctx);inputTensors.push_back(input);}else{weights.push_back(item.weights());}}OnnxAttrs attrs(node, ctx);auto name = attrs.get<std::string>("name", "");auto info = attrs.get<std::string>("info", "");// Create plugin from registryauto registry = getPluginRegistry();auto creator = registry->getPluginCreator(name.c_str(), "1", "");if(creator == nullptr){printf("%s plugin was not found in the plugin registry!", name.c_str());ASSERT(false, ErrorCode::kUNSUPPORTED_NODE);}nvinfer1::PluginFieldCollection pluginFieldCollection;pluginFieldCollection.nbFields = 0;ONNXPlugin::TRTPlugin* plugin = (ONNXPlugin::TRTPlugin*)creator->createPlugin(name.c_str(), &pluginFieldCollection);if(plugin == nullptr){LOG_ERROR(name << " plugin was not found in the plugin registry!");ASSERT(false, ErrorCode::kUNSUPPORTED_NODE);}std::vector<std::shared_ptr<TRT::Tensor>> weightTensors;for(int i = 0; i < weights.size(); ++i){auto& weight = weights[i];std::vector<int> dims(weight.shape.d, weight.shape.d + weight.shape.nbDims);auto onnx_dtype = convert_trt_datatype((::onnx::TensorProto::DataType)weight.type);if(onnx_dtype == TRT::DataType::Unknow){LOG_ERROR("unsupport weight type: " << weight.type);}std::shared_ptr<TRT::Tensor> dweight(new TRT::Tensor(dims, onnx_dtype));memcpy(dweight->cpu(), weight.values, dweight->bytes());weightTensors.push_back(dweight);}plugin->pluginInit(name, info, weightTensors);auto layer = ctx->network()->addPluginV2(inputTensors.data(), inputTensors.size(), *plugin);std::vector<TensorOrWeights> outputs;for( int i=0; i< layer->getNbOutputs(); ++i )outputs.push_back(layer->getOutput(i));return outputs;
}

從上述代碼中我們可以看到這里統一定義了一個 op_type 為 Plugin 的 op 算子，這也是為什么我們需要將 LayerNorm 的 op_type 修改為 Plugin 的原因。此外，我們會解析這個 node 的 attr 屬性拿到 name 和 info，接著我們在 createPlugin 創建插件時的名字就是使用的 name，這也就解釋了為什么我們需要添加一個 name 的屬性并命名為插件名

關于 Plugin 封裝的更多細節大家感興趣的可以看看：5.5.tensorRT基礎(2)-封裝插件過程，并實現更容易的插件開發

同樣我們可以利用 onnx_graphsurgeon 完成 ONNX 的修改，代碼如下：

import onnx
import onnx_graphsurgeon as gs# 加載 ONNX 模型
input_model_path = "rtmo-s_8xb32-600e_body7-640x640.onnx"
output_model_path = "rtmo-s_8xb32-600e_body7-640x640.plugin.onnx"
graph = gs.import_onnx(onnx.load(input_model_path))# 遍歷圖中的所有節點
for node in graph.nodes:if node.op == "LayerNormalization":node.op = "Plugin"# 添加自定義屬性node.attrs["name"] = "LayerNormPlugin"node.attrs["info"] = "This is custom LayerNormalization node"# 刪除無用的節點和張量
graph.cleanup()# 導出修改后的模型
onnx.save(gs.export_onnx(graph), output_model_path)

在終端執行該代碼后會生成 rtmo-s_8xb32-600e_body7-640x640.plugin.onnx 文件即我們修改好的 ONNX，我們一起來看看這個 ONNX 的 LayerNorm 節點與前面的又有哪些不同呢？

在這里插入圖片描述

從上圖中我們可以看到 LayerNorm 節點的 op_type 修改為了 Plugin，同時它新增了 name 屬性，名字為 LayerNormPlugin，符合我們的預期

接著我們就來看看如何利用封裝好的 plugin 來進行 engine 構建和模型推理，我們這里就直接利用 C++ API 來構建 engine 并推理，指令如下：

cd tensorRT_Pro-YOLOv8
mkdir build && cd build
cmake .. && make -j64
cd ../workspace
./pro rtmo

輸出如下所示：

在這里插入圖片描述

可以看到模型成功構建了，說明 plugin 加載成功了，不過在推理的時候出現了問題，具體錯誤信息如下：

[error][cuda_tools.cpp:25]:CUDA Runtime error cudaStreamSynchronize(stream_) # an illegal memory access was encountered, code = cudaErrorIllegalAddress
[error][cuda_tools.cpp:25]:CUDA Runtime error cudaMemcpyAsync(gpu<unsigned char>() + offset_location, src, copyed_bytes, cudaMemcpyDeviceToDevice, stream_)

從錯誤信息可以看出是 CUDA Runtime 內存訪問時出現了非發的訪問，程序在試圖訪問未分配或已經釋放的 GPU 內存地址時報錯。從封裝插件的代碼上看博主也沒有找到問題發生的原因，但是應該是封裝的 layernorm 插件出現了問題，因為未封裝的可以正常運行。

那會不會是 epsilon 和 axis 變量的解析導致的呢？感覺也不太可能，博主目前也不知道是什么原因，大家感興趣的可以幫忙看看

2.4 layerNorm plugin核函數實現淺析

那其實整個插件最重要的部分就是其核函數的實現，上面博主并沒有去分析是怎么實現的，只是簡單分享了下其使用以及使用過程中存在的一些問題，下面博主讓 ChatGPT 簡單分析了下其核函數的實現

這里先臨時補充下 LayerNormalization 的理論推導，LN 和 BN 的計算其實是差不多的，兩者都是在做 normalization 只不過是在不同的維度上，具體實現如下：

step1：mean 和 std 的計算
$\mu^l=\frac1H\sum_{i=1}^Ha_i^l\quad\sigma^l=\sqrt{\frac1H\sum_{i=1}^H(a_i^l-\mu^l)^2}$
step2：normalization
$\hat{\mathbf{a}}^{l}=\frac{\mathbf{a}^{l}-\mu^{l}}{\sqrt{(\sigma^{l})^{2}+\epsilon}}$
step3：創建一個可以學習的 mean 和 var，最后再接一個激活函數
$\mathbf{h}^{l}=f(\mathbf{g}^{l}\odot\hat{\mathbf{a}}^{l}+\mathbf{b}^{l})$
總結下來：
$\mathbf{h}=f(\frac{\mathbf{g}}{\sqrt{\sigma^2+\epsilon}}\odot(\mathbf{a}-\mu)+\mathbf{b})$
Transformer 中 LN 的處理如下圖所示：

在這里插入圖片描述

下圖更加形象的說明了 BN 和 LN 的差異:

在這里插入圖片描述

LayerNorm 核函數的代碼實現如下：

void __global__ layernorm_kernel<float>(const float* x, const float* weight, const float* bias, float* y, int N, int C, float epsilon){int idx = blockIdx.y * blockDim.y + threadIdx.y;if(idx >= N) return;// x: N, C// y: N, C// weight: C// bias:   Cconst float* px = x + idx * C;float*       py = y + idx * C;// reduce sumfloat sq = 0.0f;float s  = 0.0f;float diver = 1.0f / C;for(int ic = threadIdx.x; ic < C; ic += warpSize){float x = px[ic];s += x;sq = fmaf(x, x * diver, sq);}for (int mask = 16; mask > 0; mask /= 2)s += __shfl_xor_sync(0xffffffff, s, mask);for (int mask = 16; mask > 0; mask /= 2)sq += __shfl_xor_sync(0xffffffff, sq, mask);float mean = s / C;float rstd = rsqrtf(sq - mean * mean + epsilon);for(int ic = threadIdx.x; ic < C; ic += warpSize) py[ic] = (px[ic] - mean) * weight[ic] * rstd + bias[ic];
}

這段代碼展示了一個用于計算 Layer Normalization 的 CUDA 核函數，下面我們簡單分析下：(from ChatGPT)

1. 輸入和輸出參數

x: 輸入特征數據，形狀為 (N, C)。
weight: Layer Normalization 的權重，形狀為 (C)。
bias: Layer Normalization 的偏置，形狀為 (C)。
y: 輸出數據，形狀為 (N, C)。
N: 樣本數量。
C: 特征維度數。
epsilon: 用于穩定計算的小常數。

2. 核函數流程

索引計算：使用 blockIdx 和 threadIdx 來確定當前線程負責的樣本 idx。
早退機制：如果 idx 超出樣本數 N，則直接返回。
局部變量初始化：
- sq：用于累計平方和。
- s：用于累計和。
- diver：1/C，預計算用于優化。
并行累加：
- 通過循環每個線程處理多個特征元素，并使用步長為 warp size（通常為 32）。
- 使用 fmaf 函數累計平方和，這是一個數學加速函數，用于計算 x * x * diver + sq。
跨線程通信：
- 使用 __shfl_xor_sync 實現線程間的加法操作，合并計算結果。這是提高效率的關鍵，因為它允許在沒有全局內存訪問的情況下進行快速的數據交換。
均值和標準差計算：
- 均值 mean：s / C。
- 逆標準差 rstd：使用 rsqrtf 函數計算 (sq - mean * mean + epsilon) 的逆平方根。
輸出計算：
- 用計算出的均值和標準差來標準化輸入 x，并應用權重和偏置。

該代碼主要有以下幾點可以提高效率：

并行處理：通過循環和 warp size 來實現特征維度的并行處理。
數學加速函數 fmaf：提高浮點計算的精度和效率。
跨線程通信：使用 __shfl_xor_sync 函數減少內存訪問和加速線程間的數據交換。

這段代碼有效地實現了Layer Normalization，利用 CUDA 的并行處理和硬件加速特性優化了計算性能。

OK！以上就是關于 LayerNorm Plugin 使用和說明的全部內容了，若有問題，歡迎各位看官批評指正。

結語

博主在這里對 CUDA-BEVFusion 中的 layernorm plugin 進行了簡單的使用，并完成了插件的加載和模型推理，總的來說使用流程比較簡單（畢竟代碼都是現成的😂），將插件編譯成一個動態庫，然后交給 tensorRT 去加載就行，需要注意的是 ONNX 的一些修改。此外，plugin 的封裝部分存在的問題博主并沒有解決，歡迎大家交流討論。

通過 layernorm plugin 的使用博主又重新回顧了一遍之前學習的一些知識，雖然整個過程比較簡單，不過總歸還是有收獲的🤗