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 的使用博主又重新回顾了一遍之前学习的一些知识，虽然整个过程比较简单，不过总归还是有收获的🤗