承接《TensorRT 自定义算子插件实战》系列前两篇第一篇 customScaledTanh单输入、逐元素、带参、第二篇 customGatedTanh双输入、融合、带参。本篇实现第三个形态、也是真正拉开差距的一个customMaxpool——一个 2×2 核、stride1、无参数的窗口算子。它的难点不再是多输入而是两个全新问题算子没有参数参数机制整个被砍掉、输出尺寸发生变化不再等于输入。这三个算子加起来正好覆盖了自定义插件的三大典型形态。前两篇链接TensorRT 自定义算子插件实战一从零手写 customScaledTanhTensorRT 自定义算子插件实战二双输入融合算子 customGatedTanh目录一、算子的区别二、为什么需要自定义插件三、为什么需要 Plugin 和 PluginCreator 两个类四、本案例的算子与网络五、Python 端导出无参数的窗口算子六、C 端头文件——无参数后的减法七、C 端CUDA 核函数——真正的难点八、C 端Plugin 类的实现——两处从抄到算8.1 差异一getOutputDimensions 真正计算输出尺寸8.2 差异二serialize 空实现九、C 端PluginCreator 类的实现——被架空的工厂十、构建与验证十一、实践经验十二、小结一、算子的区别维度ScaledTanh / GatedTanh前两篇customMaxpool本篇算子类别逐元素空间邻域窗口取最大值算子参数有k/a、a/b无输出形状 输入形状≠ 输入H-1、W-1线程与数据的关系线程 index 元素 index一一对应线程 输出元素再反查输入窗口getOutputDimensionsreturn inputs[0]真正计算 H-1、W-1构造函数3 个含带参、反序列化 buffer1 个仅 nameserializememcpy mParams空实现返回 0Creator 的 mAttrs有参数空无参数可传一句话概括本篇的核心无参数让插件里参数那套机制多个构造、mParams、serialize、Creator 解析整体消失外壳变得极简输出尺寸变化让getOutputDimensions从抄输入变成真算也让 kernel 的索引从和元素一一对齐变成对齐输出、反查输入窗口。前者是减法后者是难度真正的来源。二、为什么需要自定义插件和TRT 不认识customMaxpool这个节点必须用插件实现。不同的是TRT 本身就有标准 MaxPool 层——这里写它纯粹是为了演示空间邻域算子这一类插件的写法卷积、池化、下采样、邻域统计都是这类。所以本篇的示范意义在于当你需要一个 TRT 没有的、或者你想自己控制的窗口运算时怎么写空间邻域插件。三、为什么需要 Plugin 和 PluginCreator 两个类机制与前两篇完全相同不重复。本篇要额外强调的一个现象无参数时PluginCreator 这个类虽然还在但被架空了——它的 mAttrs 是空的、createPlugin/deserializePlugin 都只是直接 new 一个 Plugin。原因不变TRT 只和 Creator 打交道Creator 是插件的注册入口即使没有参数要传递注册这件事也必须有。所以 Creator 删不掉只是参数部分归零。四、本案例的算子与网络算子定义窗口算子输出尺寸与输入不同输入 [1,2,5,5] → 2×2 步长为1的最大池化 → 输出 [1,2,4,4]池化核 2×2、stride1因此输出高度 输入高度 − 15→4宽度同理5→4。通道数不变2。网络结构单输入、单输出和前两篇对照前两篇是卷积 → 逐元素处理输出尺寸不变本篇是卷积 → 窗口池化输出在 H、W 上各小 1——这正是getOutputDimensions和 kernel 索引要处理的新问题。尺度变化如下图所示以单通道为例五、Python 端导出无参数的窗口算子和前两篇最大的区别symbolic不再接收任何标量参数g.op也不带_f属性后缀——算子完全由数据和内核行为定义。python 内容不多撰写如下importtorchimporttorch.onnximporttorch.nnasnnimportonnximportonnxsimclassCustomMaxpoolImpl(torch.autograd.Function):staticmethoddefsymbolic(g,x):# 无参数只传张量 x不带任何属性returng.op(custom::customMaxpool,x)staticmethoddefforward(ctx,x):# 与 TRT 侧 kernel 对应的行为2×2 最大池化stride1returntorch.max_pool2d(x,kernel_size2,stride1)classCustomMaxpool(nn.Module):defforward(self,x):returnCustomMaxpoolImpl.apply(x)classModel(torch.nn.Module):def__init__(self):super().__init__()self.convnn.Conv2d(1,2,(3,3),padding1)self.maxpoolCustomMaxpool()forminself.modules():ifisinstance(m,nn.Conv2d):nn.init.kaiming_normal_(m.weight,modefan_out,nonlinearityrelu)defforward(self,x):xself.conv(x)xself.maxpool(x)returnxdefexport_norm_onnx(input,model):file./sample_customMaxpool.onnxtorch.onnx.export(modelmodel,args(input,),ffile,input_names[input0],output_names[output0],opset_version11)model_onnxonnx.load(file)model_onnx,checkonnxsim.simplify(model_onnx)assertcheck onnx.save(model_onnx,file)if__name____main__:torch.manual_seed(1)inputtorch.rand(1,1,5,5)modelModel().eval()export_norm_onnx(input,model)与前两篇的差异symbolic 无属性前两篇是g.op(..., k_fk, a_fa)带标量属性本篇g.op(custom::customMaxpool, x)一个属性都不带——这是无参数在 Python 端的体现。forward 用现成算子前两篇 forward 手写公式tanh/sigmoid 组合本篇直接调torch.max_pool2d——因为 PyTorch 有这个算子forward 只要能算对即可真正实现池化逻辑的是后面 CUDA kernel。六、C 端头文件——无参数后的减法无参数最直观的体现就是头文件变短前两篇的 Plugin 有带参构造 反序列化 buffer 构造 默认构造三个本篇只剩一个 name 构造mParams 结构体整个删除。差异片段classCustomMaxpoolPlugin:publicIPluginV2DynamicExt{public:CustomMaxpoolPlugin(conststd::stringname);// 唯一的构造parse / clone / 反序列化 共用// ... 其余接口声明与前两篇完全相同 ...private:conststd::string mName;// 注意没有 mParams 成员了std::string mNamespace;};classCustomMaxpoolPluginCreator:publicIPluginCreator{// ... 与前两篇相同但 mAttrs 为空 ...};为什么无参数就能砍到只剩一个构造回顾前两篇带参构造是parse 阶段接收 k/a 并存入 mParams反序列化 buffer 构造是从引擎字节流恢复 mParams。既然没有参数mParams 不存在了这两个构造自然失去意义——parse 后无需存参数反序列化也无需恢复参数。于是三个构造收敛成一个name构造parse、clone、deserialize 三个场景复用同一个构造。七、C 端CUDA 核函数——真正的难点核函数是本篇最难、也最值得看的部分。它和逐元素算子有本质区别线程不再和输入元素一一对应而是对齐输出元素再反查输入窗口。cu 内容不多完整代码如下// custom-Maxpool.cu#includecuda_runtime.h#includemath.h#includecuda_fp16.h// 辅助函数求输入 in 在通道 c 内、以 (oy, ox) 为左上角的 2×2 窗口的最大值__device__floatWindow_Max(constfloat*in,intc,intinH,intinW,intoy,intox){floatm-3.4e38f;for(inti0;i2;i)// 列偏移 i{for(intj0;j2;j)// 行偏移 j{// 一页一页算先跳到通道 c 的页首再加行偏移再加列偏移inttemp_indexc*inH*inW(oyj)*inW(oxi);min[temp_index]m?in[temp_index]:m;}}returnm;}// 主核线程 输出元素不是输入__global__voidcustomMaxpoolKernel(constfloat*inputs,float*outputs,intc,intinH,intinW,intoutH,intoutW,constintnElements){constintindexblockIdx.x*blockDim.xthreadIdx.x;inttotalc*outH*outW;// 输出的总元素数if(indextotal)// 物理裁剪超出的线程直接退出return;// 从输出下标反解出属于哪个通道 c、哪个输出行 oy、哪个输出列 oxinttemp_cindex/(outH*outW);// 先定通道inttemp_rindex%(outH*outW);// 通道内的余数inttemp_oytemp_r/outW;// 输出行inttemp_oxtemp_r%outW;// 输出列outputs[index]Window_Max(inputs,temp_c,inH,inW,temp_oy,temp_ox);}voidcustomMaxpoolImpl(constfloat*inputs,float*outputs,constintnElements,cudaStream_t stream){dim3blockSize(256,1,1);dim3gridSize(ceil(float(nElements)/256),1,1);// 注意这里把形状写死成了 2/5/5/4/4c2, inH5, inW5, outH4, outW4见第十节坑customMaxpoolKernelgridSize,blockSize,0,stream(inputs,outputs,2,5,5,4,4,nElements);}索引逻辑是其与逐元素算子最大的分水岭。前两个算子ScaledTanh、GatedTanh里index既是输入下标也是输出下标一个线程算一个对应位置的元素。但最大池化的输入输出元素对不上输入 [1,2,5,5] 有 50 个元素输出 [1,2,4,4] 只有 32 个每个输出元素要看输入里 4 个元素。所以线程数量必须由输出决定线程 输出元素total c * outH * outWguard 用输出元素数裁剪index total就 return。这和你第一篇看过的逐元素 guard含义完全不同——那里 index 是输入下标这里 index 是输出下标。从输出下标反解坐标一个输出元素 index 要回答它在哪个通道、哪一行、哪一列用一连串整除/取模拆出来index ÷ (outH×outW) → 通道 c index % (outH×outW) → 通道内线性位置再 ÷ outW → 行 oy再 % outW → 列 ox反查输入窗口有了 (c, oy, ox)去输入里取以 (oy, ox) 为左上角的 2×2 邻域窗口内元素下标是c*inH*inW (oyj)*inW (oxi)j 是行内偏移 0/1i 是列内偏移 0/1。这是逐元素算子里绝不会出现的空间邻域寻址。一句话逐元素算子是一个线程对一个元素窗口算子是一个线程对一个输出元素、但要从输入的邻域里取数据。索引的对齐对象从输入变成了输出这就是空间邻域算子难的地方。无参数在 kernel 里的体现kernel 签名里没有 a/b 这类标量参数只剩数据指针、形状参数和元素数——因为算子行为取 2×2 最大值是写死的不需要任何可调参数。八、C 端Plugin 类的实现——两处从抄到算无参数、输出尺寸变化带来的改动集中在两个方法其余接口构造、getPluginType、clone、destroy……与前两篇逐字相同。8.1 差异一getOutputDimensions 真正计算输出尺寸前两篇直接return inputs[0]逐元素算子输出 输入。最大池化输出在 H、W 上各小 1必须真算DimsExprsCustomMaxpoolPlugin::getOutputDimensions(int32_toutputIndex,constDimsExprs*inputs,int32_tnbInputs,IExprBuilderexprBuilder)noexcept{DimsExprs outinputs[0];// 先拷贝N、Cd[0]、d[1]不变// NCHWd[2]H、d[3]W输出 输入 - 1因为 2×2 核、stride1out.d[2]exprBuilder.operation(DimensionOperation::kSUB,*inputs[0].d[2],*exprBuilder.constant(1));// H-1out.d[3]exprBuilder.operation(DimensionOperation::kSUB,*inputs[0].d[3],*exprBuilder.constant(1));// W-1returnout;}动态维度要用 IExprBuilder 运算因为插件是 DynamicExt 版本H、W 可能是动态的不能用普通整数相减要用exprBuilder.operation(kSUB, d[2], constant(1))声明输出维度 输入维度 − 1。这样 TRT 就知道输出是[N, C, H-1, W-1][1,2,5,5]进来给到[1,2,4,4]。8.2 差异二serialize 空实现无参数 → 没有东西可序列化size_tCustomMaxpoolPlugin::getSerializationSize()constnoexcept{return0;}voidCustomMaxpoolPlugin::serialize(void*buffer)constnoexcept{// 无参数什么都不写return;}对照前两篇的memcpy(buffer, mParams, sizeof(mParams))——mParams 没了序列化也随之归零。九、C 端PluginCreator 类的实现——被架空的工厂无参数让 Creator 也大幅变空mAttrs 不再注册任何 PluginFieldcreatePlugin/deserializePlugin 直接 new。CustomMaxpoolPluginCreator::CustomMaxpoolPluginCreator(){// 无参数不再 emplace_back 任何 PluginField// mAttrs.emplace_back(PluginField(a, ...)); ← 已注释掉// mFC.nbFields 0;}IPluginV2*CustomMaxpoolPluginCreator::createPlugin(constchar*name,constPluginFieldCollection*fc)noexcept{// 无参数不需要从 fc 里解析任何值returnnewCustomMaxpoolPlugin(name);}IPluginV2*CustomMaxpoolPluginCreator::deserializePlugin(constchar*name,constvoid*serialData,size_t serialLength)noexcept{// 无参数反序列化也不需要恢复任何值returnnewCustomMaxpoolPlugin(name);}为什么 Creator 还是不能删因为REGISTER_TENSORRT_PLUGIN(CustomMaxpoolPluginCreator)注册的是这个 CreatorTRT 靠它按 op_type domain 找到如何创建这个插件。即使它内部什么都不做注册这个入口本身不能省。所以无参数时 Creator 是空壳但必须存在。十、构建与验证main函数比较简单直接读取相关的onnx文件进行本地引擎构建再推理即可。与上篇的唯一区别就是推理时候读取的onnx文件不一样。#includeiostream#includememory#includeutils.hpp#includemodel.hppusingnamespacestd;intmain(intargc,charconst*argv[]){Modelmodel(models/onnx/sample_customMaxpool.onnx,Model::precision::FP16);if(!model.build()){LOGE(fail in building model);return0;}if(!model.infer()){LOGE(fail in infering model);return0;}return0;}验证方式相同PyTorch 跑 ONNX、C 加载 TRT 引擎跑插件对比输出。我的实现里两者完全一致证明输出尺寸变化 → getOutputDimensions 算对 → kernel 对齐输出反查输入窗口这条链路成立。python程序的输出结果为cpp程序的输出结果为实验发现Python 与 C 的推理结果完全一致基本可以确定软件没有问题。十一、实践经验host 封装把形状硬编码了。customMaxpoolImpl里customMaxpoolKernel...(inputs, outputs, 2, 5, 5, 4, 4, ...)把 c/inH/inW/outH/outW 写死成当前模型的尺寸。这能让 demo 跑通但换了输入尺寸就错。正确做法应从inputDesc[0].dims或 configurePlugin 保存的尺寸动态取出 H、W算出 outH、outW 再传 kernel。enqueue 传输入元素数、kernel 用输出元素数口径不一致。enqueue 里 nElements 输入总元素数2×5×550而 kernel 里 guard 用的是输出 total2×4×432。本例因为输出 输入guard 正确裁剪了但如果输入输出关系反过来或开方不整grid 划分和 guard 就可能出问题。栅格划分应基于实际要执行的输出元素数。线程对齐对象变了。逐元素算子线程对齐输入元素窗口算子线程必须对齐输出元素因为输出个数决定要算几个数。很多人第一次写池化 kernel下意识按输入去分线程结果边界全错——先想清楚线程数该由谁决定。动态维度运算要用 IExprBuilder。DynamicExt 插件里算输出尺寸不能直接写d[2] - 1必须用exprBuilder.operation(kSUB, ...)否则动态 shape下会错。对于以上第一点为了达到更好的兼容性customMaxpoolImpl函数内部可以加上如下代码// 从输出取 outH、outW从输入取 c、inH、inWNCHWd[0]N、d[1]C、d[2]H、d[3]WintcinputDesc[0].dims.d[1];intinHinputDesc[0].dims.d[2];intinWinputDesc[0].dims.d[3];intoutHoutputDesc[0].dims.d[2];intoutWoutputDesc[0].dims.d[3];// 再传进 customMaxpoolImpl → kernel对于以上第二点所以enqueue()函数更好的实现方法是用outputDesc来算元素个数int32_tCustomMaxpoolPlugin::enqueue(constPluginTensorDesc*inputDesc,constPluginTensorDesc*outputDesc,constvoid*const*inputs,void*const*outputs,void*workspace,cudaStream_t stream)noexcept{/* * Plugin的核心的地方。每个插件都有一个自己的定制方案 * Plugin直接调用kernel的地方 */intnElements1;for(inti0;ioutputDesc[0].dims.nbDims;i){nElements*outputDesc[0].dims.d[i];}customMaxpoolImpl(static_castconstfloat*(inputs[0]),static_castfloat*(outputs[0]),nElements,stream);return0;}另外两点已经在代码中有所体现不再赘述。十二、小结三个算子写到这里正好覆盖自定义插件的三大形态做一个收尾对照形态代表算子实现难点单输入 · 逐元素 · 带参customScaledTanh参数在onnx 属性 → mFC → mParams → 序列化间传递双输入 · 融合 · 带参customGatedTanhenqueue 多取一个指针、supports 多一个 case单输入 · 窗口 · 无参customMaxpool无参数机制砍掉输出尺寸变化线程对齐输出、反查输入窗口无参数让插件结构减到最简一个构造、空 serialize、空 Creator你因此看清了哪些外壳是参数机制撑起来的、哪些是注册机制必需的输出尺寸变化则逼你把getOutputDimensions从抄输入升级成真算把 kernel 索引从元素对齐升级成对齐输出、空间邻域寻址。这三篇加起来自定义插件里外壳与内核的开发范式就完整了。