4W字长文 | Sparse4Dv3的TensorRT部署调优指南~

作者 | Thomas Wilson  编辑 | 自动驾驶Daily

原文链接：https://zhuanlan.zhihu.com/p/715179777

点击下方 卡片 ，关注“ 自动驾驶之心 ”公众号

戳我-> 领取 自动驾驶近15个 方向 学习 路线

>> 点击进入→ 自动驾驶之心 『 模型部署 』 技术交流群

本文只做学术分享，如有侵权，联系删文

Brief

本文分享主题： 如何在个人工作站以及车载NVIDIA ORIN 上部署Sparse4Dv3端到端感知方案。 Sparse4D是基于稀疏Transformer范式的高性能高效率的长时序融合的感知算法。截止2024年9月，该方案在nuscenes纯视觉榜单上以mAP=0.668排名位列第一：

以下为算法的通用架构图[1]：

当前，我测试了该算法在多种数据集上的泛化效果，感觉确实还不错。基于此，我重构了源代码，主要工作：剔除了源代码对MMDetection3D和mmcv-full的依赖，这使得该算法更加轻量化，无论是训练、推理还是部署都变的更加友好。目前，该仓库测试下来，无论是在个人的工作站还是远程集群环境（这里主要测试了NVIDIA Ampere架构系列显卡、NVIDIA Volta架构系列显卡还有Hopper架构显卡），其安装简单无复杂依赖（如需升级，只要更新CUDA和torch相关组件版本即可），轻松完成大规模训练和推理任务。

其实对于大部分自动驾驶量产公司来说，当模型的算法方案确定后，数据的采集到标注爬坡是需要一定周期的，日常工作除了数据的处理和模型调优迭代外，大部分精力需要投入到模型的部署、模型的高性能推理和模型的车端C++代码开发，毕竟能上车才是王道。

本人在对Sparse4Dv3模型部署和C++代码开发过程中，感觉并不是一番风顺。因此，我想借这篇文章将个人踩过的所有坑和经验分享给大家，希望大家在使用过程中尽量少走弯路。号外号外，感兴趣的朋友可以clone我的 GitHub Repo : SparseEnd2End。部署代码请大家关注仓库中的deploy文件夹，该代码目前已经开源。车端C++代码在持续更新填充中，请大家关注onboard文件夹。好记性，不如烂笔头，本文会长期更新，内容将会涉及后续的调优以及持续加速......

github.com/ThomasVonWu/SparseEnd2End

大家在使用该仓库过程中如果遇到什么问题也请分享给我，共同学习进步嘛。

在阅读本文之前，需具备以下基础知识和相关工具使用经验，这将方便大家快速理解本篇文章的核心内容（抱歉，时间有限，后续涉及基础工具的使用方法我就不详细展开了）。另外，本篇文章的部署方案无需依赖其他部署仓库，如：MMdeploy, etc：

1. PyTorch模型导出ONNX中间格式的方法，熟悉可视化ONNX节点工具: netron的使用；
2. 推理引擎 ONNX Runtime 和 TensorRT的安装及使用方法，包括：TensorRT python API 和 C++ API使用；
3. 熟悉TensorRT 工具Polygraphy的python API使用方法和基本的命令行调用指令;
4. 基本的CUDA编程知识：核函数的编写与启动，常用的内存模型：全局内存、共享内存，etc；
5. Makefile编程语法，C++编译规则和nvcc 编译cuda程序规则；

通过这篇文章你将了解并学习到以下进阶知识：

1. 模型转换完，如何确保数据流的无损传递以及推理结果的一致性验证方法，即：如何debug问题和校验结果准确性 (我想这是大家最关心的部分，毕竟网上大几百的课程也未必会详细的告诉大家这些方法)；
2. 如何在ONNX注册一个自定义算子？紧接着，如何将该自定义算子注册为TensorRT Plugin，最后你一定关心：如何将该自定义plugin加载到当前模型的engine中，并使用TensorRT python API完成python脚本推理链路和TesnsorRT C++ API完成车端推理链路;

好的，一切就绪，现在让我们愉快的开始吧～

Sparse4Dv3 Model Deployment Pipeline

本文配置的部署环境以及使用的工具版本，如下：

====================================================================================================================
||  Config Environment Below:
||  UBUNTU              : 20.04
||  TensorRT LIB        : /mnt/env/tensorrt/TensorRT-8.5.1.7/lib
||  TensorRT INC        : /mnt/env/tensorrt/TensorRT-8.5.1.7/include
||  TensorRT BIN        : /mnt/env/tensorrt/TensorRT-8.5.1.7/bin
||  CUDA_LIB            : /usr/local/cuda-11.6/lib64
||  CUDA_ INC           : /usr/local/cuda-11.6/include
||  CUDA_BIN            : /usr/local/cuda-11.6/bin
||  CUDNN_LIB           : /mnt/env/tensorrt/cudnn-linux-x86_64-8.6.0.163_cuda11-archive/lib
||  CUDASM              : sm_86
||  PYTORCH             : 1.13.0
||  ONNX                : 1.14.1
||  ONNXRUNTIME         : 1.15.0
||  ONNXSIM             : 0.4.33
||  CUDA-PYTHON         : 1.15.0
||  NETRON              : 7.7.8
||  POLYGRAPHY          : 0.49.9
====================================================================================================================

部署的大体思路如下：

1. Sparse4Dv3模型部署的过程需要将我们将训练好的模型 .pth 文件，转换为中间文件 .onnx，最后转换为*.engine 文件。该过程需要解决PyTorch模型与ONNX框架算子的兼容性以及模型运行加速两大需求; 模型pth下载链接：https://drive.google.com/file/d/1sSMNB7T7LPKSr8nD9S_tSiu1mJrFMZ1I/view?usp=sharing**
2. Sparse4Dv3 PyTorch模型实际上就是一个计算图。模型部署时通常需要我们将模型转换成静态的计算图，即没有控制流（分支语句、循环语句）的计算图，这点很重要，Sparse4Dv3 Head模型转换过程我们就会遇到这个问题；
3. PyTorch 框架自带对 ONNX 的支持，只需要我们构造一组随机的输入，并对模型调用 torch.onnx.export 即可，完成 PyTorch 到 ONNX 的转换。
4. * 推理引擎 ONNX Runtime 对 ONNX 模型有原生的支持，提供了python API 和 C++ API。给定一个 . onnx 文件，只需简单使用 ONNX Runtime 的 Python API 就可以完成模型推理。依据上述工具的使用，我们可以完成PyTorch到ONNX的推理一致性验证。
5. 推理引擎 TensorRT 提供了Python API 和 C++ API。给定一个 .engine 文件，只需要简单使用 TensorRT 的 Python API 就可以在python脚本中完成模型推理。依据上述构造过程，我们可以完成PyTorch到ONNX的推理一致性验证。最后，通过调用C++ API就可以完成代码在车端仓库的部署了(简单点说，就是将TensorRT的python API做一次C++ API接口映射即可)。 *

部署工作开始前，我们首先分析下Sparse4Dv3 PyTorch模型结构，这里以配置：输入img_size : 256x704, 模型backbone : Resnet50, 模型精度Precision : fp32为例，模型结构大体可以拆分成两个部分：imgBackbone和sparseTrans formerHead。

1）imgBackbone由Resnet50+FPN组成，详细结构如下：

(imgBackbone): ResNet(
    (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu): ReLU(inplace=True)
    (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
    (layer1): ResLayer(
      (0): Bottleneck(
        (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): Bottleneck(...)
      ...
    )
    ***********************************************************Resnet50重复堆叠的模块我这里省略了哈
  )
  init_cfg={'type'




    
: 'Pretrained', 'checkpoint': 'ckpt/resnet50-19c8e357.pth'}
  (img_neck): FPN(
    (lateral_convs): ModuleList(...)
    (fpn_convs): ModuleList(
      (0): ConvModule(...)
  )
  init_cfg={'type': 'Xavier', 'layer': 'Conv2d', 'distribution': 'uniform'})

2）sparseTrans formerHead 由 anchorencoder + 39个Op构成的ModuleList组成，其中，不包含instance_bank layer, loss_cls, loss_reg, loss_box, loss_cns, etc. 原因如下：

• instance_bank layer: 该模块主要功能，缓存历史帧instance实例，并将历史帧anchor投影到当前帧，最后更新trackid和confidence，而且，该模块并未包含有效的nn.Module部分(这里有效指的是训练过程并没有需要更新梯度的tensor变量，也就不涉及到权重和偏置组成部分)，所以为了成功导出ONNX，后续会将该部分从head剥离开；
• 模型部署的数据流走的是推理链路，所以和训练相关的模块需要全都丢弃，这里包含：loss_cls, loss_reg, loss_box, loss_cns, depth_branch, grid_mask, etc；

(sparseTransformerHead): Sparse4DHead(
    (anchor_encoder): SparseBox3DEncoder(
      (pos_fc): Sequential(
        (0): Linear(in_features=3, out_features=128, bias=True)
        (1): ReLU(inplace=True)
        (2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
        (3): Linear(in_features=128, out_features=128, bias=True)
        (4): ReLU(inplace=True)
        (5): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
        (6): Linear(in_features=128, out_features=128, bias=True)
        (7): ReLU(inplace=True)
        (8): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
        (9): Linear(in_features=128, out_features=128, bias=True)
        (10): ReLU(inplace=True)
        (11): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
      )
*********************************************************size_fc/yaw_fc/vel_fc我省略了哈 
    )
    (layers): ModuleList(
      (0): DeformableAttentionAggr(...)
      (1): AsymmetricFFN(...)
      (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
      (3): SparseBox3DRefinementModule(...)
        (cls_layers): Sequential(...)
        (quality_layers): Sequential(...)
      (4): MultiheadAttention(...)
      (5): MultiheadAttention(...)
      ***********************************************************head重复的op我也省略了哈  
    )
    (fc_before): Linear(in_features=256, out_features=512, bias=False)
    (fc_after): Linear(in_features=512, out_features=256, bias=False))

Sparse4Dv3 SparseTransfomerHead Deployment Solution

核心内容 ，之所以先介绍head部署方案，一方面，因为head是算法的核心部分，和业务需求紧密联系，另一方面。head结构复杂，部署的主要工作量体现在这一部分。相比较imgBackbone，推理前向过程需要我们很清晰的掌握head模块的数据流向以及算法机理。

一)DFA自定义CUDA算子的部署方案

首先head在转换ONNX中间格式时，我们不难发现，其中存在自定义算子DFA( Deformable Feature Aggregation )，其调用函数为deformable_aggregation_function()。很明显，该算子借用pybind11构建共享库完成了C++和CUDA扩展的自定义算子，从而解决PyTorch原算子并行问题，进而减小训练/推理显存，加速训练/推理速度。PyTorch的C++和CUDA扩展自定义算子一般流程如下：

• 第一步：使用C++编写算子的forward函数和backward函数;
• 第二步：将该算子的forward函数和backward函数使用pybind11绑定到python上;
• 第三步：使用setuptools/JIT/CMake编译打包C++工程为*.so文件;
• 第四步：在python端继承PyTorch的torch.autograd.Function类，实现静态函数forward和backward，并调用上述过程生成的动态库*.so;

Q∶什么是DFA算子，它又具备什么样的功能?论文中这样写道：[2]

DFA算子实现流程：

• 第一步：首先，将900个query instance 的 13（6KFP+7KLP）个3D关键点全部投影current timestamp的FeatureMap上，接下来通过双线性插值方法在MultiView FeatureMap上进行特征采样;
• 第二步：紧接着，在MultiView的不同的MultiScale层（即不同分辨率的特征图层）上重复执行双线性插值操作，以捕获从粗粒度到细粒度的特征特区，这有助于模型在不同尺度上理解物体的结构和细节;
• 第三步：最后，网络会使用预测的权重（通过线性层计算得到）进行加权，完成特征的聚合。这确保算法具备了：根据检测任务的贡献大小，对不同视角和尺度的特征进行合适的特征融合;

P.S.其中KFP为固定关键点，KLP为可学习关键点，文中DeformableAttention Aggr,DeformableFeatureAggregation, DFA都指的同一个自定义算子

Q:为什么需要对DFA算子做CUDA加速?

HBM( High Bandwidth Memory )[3]高带宽存储，是一种常用显存介质。顾名思义，这个存储介质有着"High Bandwidth"。在多视角多尺度特征聚合过程中，DFAOp涉及了多次HBMIO过程，因此，存储了大量的临时变量。训练过程梯度回传以及推理过程频繁的HBM访问，一方面，其占用了大量显存。另一方面，频繁IO降低了推理速度。原码中调用的是torch.nn.functional. ample[4] 接口，由于PyTorch自带的该算子底层没有做并行运算的加速优化，所以，在训练或推理过程中，我发现3090的显卡都根本吃不消，显存直接爆炸了，核心代码调用如下：

for fm in feature_maps:
    features.append(nn.functional.grid_sample(fm.flatten(end_dim=1), points_2d))
    features = torch.stack(features, dim=1)
    features = features.reshape(
        bs, num_cams, num_levels, -1, num_anchor, num_pts).permute(0, 4, 1, 2, 5, 3)

针对上述问题，作者提出了加速的需求，并给出了加速方案：

DFA算法加速流程：

• 第一步：将基本可变形聚合中的MultiScale/MultiView维度输入，同可学习的权重系数加权操作，作为CUDA线程的原子操作;
• 第二步：在K×C的维度上分配线程，实现K, C维度的完全的并行化计算;

核心代码见位置：modules/ops/src/deformable_aggregation_cuda.cu，摘取部分如下：

__global__ void deformable_aggregation_kernel(
    const int num_kernels,
    float* output,
    const float* mc_ms_feat,
    const int* spatial_shape,
    const int* scale_start_index,
    const float* sample_location,
    const float* weights,
    int batch_size,
    int num_cams,
    int num_feat,
    int num_embeds,
    int num_scale,
    int num_anchors,
    int num_pts,
    int num_groups
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= num_kernels) return;

    const float weight = *(weights + idx / (num_embeds / num_groups));
    const int channel_index = idx % num_embeds;
    idx /= num_embeds;
    const int scale_index = idx % num_scale;
    idx /= num_scale;

    const int cam_index = idx % num_cams;
    idx /= num_cams;
    const int pts_index = idx % num_pts;
    idx /= num_pts;

    int anchor_index = idx % num_anchors;
    idx /= num_anchors;
    const int batch_index = idx % batch_size;
    idx /= batch_size;

    anchor_index = batch_index * num_anchors + anchor_index;
    const int loc_offset = ((anchor_index * num_pts + pts_index) * num_cams + cam_index) <
    const float loc_w = sample_location[loc_offset];
    if (loc_w <= 0 || loc_w >= 1) return;
    const float loc_h = sample_location[loc_offset + 1];
    if (loc_h <= 0 || loc_h >= 1) return;
    
    int cam_scale_index = cam_index * num_scale + scale_index;
    const int value_offset = (batch_index * num_feat + scale_start_index[cam_scale_index]) * num_embeds + channel_index;

    cam_scale_index = cam_scale_index <    const int h = spatial_shape[cam_scale_index];
    const int w = spatial_shape[cam_scale_index + 1];

    const float h_im = loc_h * h - 0.5;
    const float w_im = loc_w * w - 0.5;

    atomicAdd(
        output + anchor_index * num_embeds + channel_index,
        bilinear_sampling(mc_ms_feat, h, w, num_embeds, h_im, w_im, value_offset) * weight
    );
}

由于单个线程的计算负载是 N X S ，而每个点最多投影到两个视角，因此，计算复杂度最多2S。到这里，我们知道了为什么需要对DFA算子做CUDA加速以及如何进行加速，那这些原理对实际部署又有什么影响呢？根据作者的加速方案，其实，我们可以进一步的优化该算子，比如，减少HBMIO次数、引入share_memory、使用半精度CUDA fp16等方法。以上内容将会在我的下一篇文章中做详细介绍，大家感兴趣的可以关注我的仓库日志更新哦～

P.S.其中N为使用的环视相机的数量，S为多尺度的层数，K为关键点的数量

Q:DFA算子怎么部署呢?具体开发流程是什么?

这里先说结论：

1. 首先，在PyTorch代码中使用符号函数：symbolic：构建DFA自定义算子PyTorch 到 ONNX 映射规则;
2. 其次，将带有自定义算子DFA，PyTorch 到 ONNX 映射规则的*.pth模型文件转化为ONNX中间文件;
3. 最后，基于指定版本的TensorRT C++ 注册plugin的 API，构建自定义DFA plugin，构建编译规则，编译该plugin，在本地将会生成动态库*.so文件;
4. 对比验证PyTorch推理结果和TensorRT推理结果，确保结果的一致性;
5. 通过trtexec command line 加载上述动态库到目标trt engine中，生成带有plugin的目标模型engine;

这里，不得不提下我之前幻想能够实现部署目标的几个解决方案的思路，以及后续为什么我最终敲定了上述方案：

被我舍弃的思路一： 拆分法

由于DFA是自定义算子，这意味着：PyTorch的ATen[5]基本库没有官方实现、ONNX节点Op不支持该算子、TensorRT(包含官方实现的Plugin)不支持该算子。因此，在转换ONNX过程必定面临失败，更别提转换TensorRT Engine了。那么比较容易想到的方法是：在自定义算子 出现的位置，我们把模型逐个拆分开，最后，在C++车端代码中逐个把模型串联起来不就可以了？一般来说这种方案没啥问题，但在Sparse4Dv3中，我们不推荐这么做，请看下面：

很显然，DFAOp单帧推理出现了6次，若要拆分模型，单单Head部分就会1拆7，这不是越搞越复杂吗？？思路一显然太笨不合适！

被我舍弃的思路二： 集成Custom DFA Plugin放入TensorRT OSS整编

TensorRT OSS[6](TensorRT Open Source Software) 是NVIDIA 高性能推理SDK仓库。我们先指定TensorRT版本，然后clone下来，在plugin文件我们可以找到官方实现的plugin。到了这里，很多博客还有视频教程会告诉我们，参考官方自定义Plugin的方法实现自定义Plugin，然后在对应的API中注册新的算子名，最后，整体编译生成新的libnvinfer.so替换原始动态库完成自定义plugin的注册。听起来逻辑清晰，好像没有什么问题，但是，实际操作下来发现问题太多了：

• 辛辛苦苦的把DFA算子按规则注册好了，准备开始编译，紧接着，各种依赖问题出现。解决了一上午，诶？还是有环境问题，这时候的你冷静思考了一下：好像关于算子的开发的工作，以及推理正确与否还没开始验证呀？我一上午在干啥呢？
• 下午的你有些焦虑，先打开google ，把问题搜索一下？额，什么都搜不到，心想：全世界就我一个人遇到了这个问题吗，不是吧？？？忍不了了，于是你开始在issue上提问，同时你打开google翻译，精心准备好问题，并确定没有语法错误，激动的把问题提交去了，坐等仓库维护者帮忙解答。不知不觉过了几天，你很烦躁，心想：他们是太忙了吗？还是有时差呀？又或者放假了？怎么现在还没有回复我的问题，怎么办？我再试试自己解决一下，突然一条消息弹出，issue有更新啦，但你赫然看到的是那句熟悉又恼火的答复：How about using the latest version TensorRT?

哈哈，这个情况不知道大家遇到的多不多，我反正经常遇到。即使这些问题在本地工作站解决了，NVIDIA ORIN 部署还是有隐患的。即使，我们在本地编译通过或者升级TensorRT完美避开旧板本的bug，最终生成新的libnvinfer.so，可是，这玩意部署到NVDIA ORIN上也用不了呀，架构不同！要么，我们在本地借用交叉编译工具完成基于ARM的共享库生成。要么，我们直接上ORIN开发板编译。且不说编译速度慢吧，谁能确保不会遇到其他环境问题？这工作量目测是个无底洞，而关于plugin开发验证工作我们实际还未开展起来，显然侧重点有问题呀！抱歉，这个思路也不适合我。

被我舍弃的思路三： 借用其他开源的仓库

借用仓库，如：tensorRT_Pro | tensorrtx | torch2trt 等。一方面需要了解并学习其使用方法。另外，也不能保证这些仓库的维护者一直持续更新，及时修复里面的BUG。最重要的是，自动驾驶使用的相关自定义算子，上述仓库大概率也没有呀（大部分是关于2D目标检测、目标分类和分割相关的算法部署方案）。所以说，准备工作挺耗费时间的，性价比不高。

基于上面的思考，最适合的方案还是一开始给出的结论最靠谱。言归正传，

第一步：建立PyTorch算子与ONNX节点映射关系：

第二步：生成ONNX中间文件。有了映射关系，我们就可以生成ONNX文件了，netron可视化如下：

这里的ONNX文件是无法成功转换trt engine，因为，ONNX中的DeformableAttentionAggrPlugin节点只是一个接口，没有对应的TensorRT Op实现的。因此，下一步需要在TensorRT中注册该算子，并将核心实现填充进去。

第三步：Custom Operator Plugin: DeformableAttentionAggrPlugin注册(我用最简单的话来解释)：

1. 打开浏览器搜索github 仓库 TensorRT，指定版本8.5，打开plugin文件夹，这里以官方的实现的第一个plugin: batchTilePlugin为例:

我们需要构建四个文件分别为：

• deformableAttentionAggrPlugin.h
• deformableAttentionAggrPlugin.cpp
• deformableAttentionAggrPlugin.cu
• Makefile

P.S.其中我们会多一个文件: deformable Attention AggrPlugin.cu，这个文件为CUDA实现的核心代码，其实和Plugin注册没什么关系，是CustomOperator :DFA实现的逻辑代码。至于使用Make file还是CMakeLists编译，依据习惯就可，无特别要求

先说第一个文件，头文件一般用来定义接口，注册自定义plugin需要构建三部分：

• part1 : 注册Plugin名称DeformableAttentionAggrPlugin，它需和ONNX节点保持一致

static const char* PLUGIN_NAME{"DeformableAttentionAggrPlugin"};
static const char* PLUGIN_VERSION{"1"};

• part2 : 定义Plugin类：DeformableAttentionAggrPlugin，这里包含了核心实现函数enqueue， 除了这个函数需要我们定义实现，其他函数基本都是依据模板直接套用（是不是很简单） ，DeformableAttentionAggrPlugin类具体又可以划分三部分组成，如下：

/// @brief PART1: Custom Plugin Class: DeformableAttentionAggrPlugin -> nvinfer1::IPluginV2DynamicExt Methods
    /*
     * clone()
     * getOutputDimensions()
     * supportsFormatCombination()
     * configurePlugin()
     * getWorkspaceSize()
     * enqueue() /// 它是核心哦！
     * attachToContext()
     * detachFromContext()
     */

/// @brief PART2: Custom Plugin Class: DeformableAttentionAggrPlugin -> nvinfer1::IPluginV2Ext Methods
    /*
     * getOutputDataType()
     */

/// @brief PART3: Custom Plugin Class: DeformableAttentionAggrPlugin -> nvinfer1::IPluginV2 Methods
    /*
     * getPluginType()
     * getPluginVersion()
     * getNbOutputs()
     * initialize()
     * getSerializationSize()
     * serialize()
     * destroy()
     * terminate()
     * setPluginNamespace()
     * getPluginNamespace()
     */

• part3 : 定义Plugin创建类：DeformableAttentionAggrPluginCreator，这部分没什么好说的，直接套用模板，改改变量名称就可以了：

/// @brief Second define a PluginCreator Class：DeformableAttenionAggrPluginCreator -> IPluginCreator
/*
 * DeformableAttentionAggrPluginCreator()
 * ~DeformableAttentionAggrPluginCreator()
 * getPluginName()
 * getPluginVersion()
 * getFieldNames()
 * createPlugin()
 * deserializePlugin()
 * setPluginNamespace()
 * getPluginNamespace()
 */

完成头文件定义，源文件实现基本都可以复用batchTilePlugin的实现，根据需要更新对应的变量名和使用即可。而主要的工作量，需要我们完成核心函数接口DeformableAttentionAggrPlugin::enqueue()的实现，而加速并行CUDA代码在*.cu实现即可(接口函数：thomas_deform_attn_cuda_forward())，完成这一步可以说大功告成了！

int32_t DeformableAttentionAggrPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
                                               const nvinfer1::PluginTensorDesc* outputDesc,
                                               const void* const* inputs,
                                               void* const* outputs,
                                               void* workspace,
                                               cudaStream_t stream) noexcept
{
    int32_t const batch = inputDesc[0].dims.d[0];
    int32_t spatial_size = inputDesc[0].dims.d[1];
    int32_t channels = inputDesc[0].dims.d[2];
    int32_t num_cams = inputDesc[1].dims.d[0];
    int32_t num_levels = inputDesc[1].dims.d[1];
    int32_t num_query = inputDesc[3].dims.d[1];
    int32_t num_point = inputDesc[3].dims.d[2];
    int32_t num_groups = inputDesc[4].dims.d[5];
    int32_t rc = 0;

    const float* value = static_castfloat*>(inputs[0]);
    const int32_t* spatialShapes = static_cast(inputs[1]);
    const int32_t* levelStartIndex = static_cast(inputs[2]);
    const float* samplingLoc = static_castfloat*>(inputs[3]);
    const float* attnWeight = static_castfloat*>(inputs[4]);

    float* output = static_cast<float*>(outputs[0]);

    rc = thomas_deform_attn_cuda_forward(stream,
                                        value,
                                        spatialShapes,
                                        levelStartIndex,
                                        samplingLoc,
                                        attnWeight,
                                        output,
                                        batch,
                                        spatial_size,
                                        channels,
                                        num_cams,
                                        num_levels,
                                        num_query,
                                        num_point,
                                        num_groups);

    return rc;
}

第四步：DFA Operators PyTorch vs. TensorRT 推理一致性验证:

代码详情见：deploy/dfa_plugin/unit_test/deformable_feature_aggregation_infer-consistency-val_pytorch_vs_trt_unit_test.py。在构建推理一致性代码前，我们需要加载Plugin动态库：这里，我们import ctyp，调用方法为 ctypes.cdll.LoadLibrary :：

ctypes.cdll.LoadLibrary(soFile)
def getPlugin(plugin_name) -> trt.tensorrt.IPluginV2:
    for i, c in enumerate(trt.get_plugin_registry().plugin_creator_list):
        logger.debug(f"We have plugin{i} : {c.name}")
        if c.name == plugin_name:
            return c.create_plugin(c.name, trt.PluginFieldCollection([]))

有了动态库文件，推理链路的开发，需要熟悉指定版本的TensorRT Python API接口的使用方法，不同版本本接口略有不同，v8.4和v8.3基本一致，而v8.5和v8.6基本一致。以v8.4(old)版本为例，举几个常用的方法：

ICuda.Engin.e类类常用方法：

I EaecutionConteat类常用方法：

CoreCodeBelow:

def inference(
    feature: np.ndarray,
    spatial_shapes: np.ndarray,
    level_start_index: np.ndarray,
    instance_feature: np.ndarray,
    anchor: np.ndarray,
    time_interval: np.ndarray,
    image_wh: np.ndarray,
    lidar2img: np.ndarray,
    engine: str,
    trt_old: bool,
    logger,
):
    bufferH = []
    bufferH.append(feature)
    bufferH.append(spatial_shapes)
    bufferH.append(level_start_index)
    bufferH.append(instance_feature)
    bufferH.append(anchor)
    bufferH.append(time_interval)
    bufferH.append(image_wh)
    bufferH.append(lidar2img)

    if trt_old:
        nIO = engine.num_bindings
        lTensorName = [engine.get_binding_name(i) for i in range(nIO)]
        nInput = sum([engine.binding_is_input(lTensorName[i]) for i in range(nIO)])

        for i in range(nInput, nIO):
            bufferH.append(
                np.zeros(
                    engine.get_binding_shape(lTensorName[i]),
                    dtype=trt.nptype(engine.get_binding_dtype(lTensorName[i])),
                )
            )

        for i in range(nInput):
            logger.debug(
                f"LoadEngine: Input{i}={lTensorName[i]}:\tshape:{engine.get_binding_shape}\ttype:{str(trt.nptype(engine.get_binding_dtype))} ."
            )
        for i in range(nInput, nIO):
            logger.debug(
                f"LoadEngine: Output{i}={lTensorName[i]}:\tshape:{engine.get_binding_shape}\ttype:{str(trt.nptype(engine.get_binding_dtype))} ."