作者 | Thomas Wilson 编辑 | 自动驾驶Daily
原文链接:https://zhuanlan.zhihu.com/p/715179777
点击下方
卡片
,关注“
自动驾驶之心
”公众号
>>
点击进入→
自动驾驶之心
『
模型部署
』
技术交流群
本文只做学术分享,如有侵权,联系删文
Brief
本文分享主题:
如何在个人工作站以及车载NVIDIA ORIN 上部署Sparse4Dv3端到端感知方案。
Sparse4D是基于稀疏Transformer范式的高性能高效率的长时序融合的感知算法。截止2024年9月,该方案在nuscenes纯视觉榜单上以mAP=0.668排名位列第一:
图一:nuscenes纯视觉榜单
以下为算法的通用架构图[1]:
图二:基于时序融合稀疏Transformer范式算法算法架构
当前,我测试了该算法在多种数据集上的泛化效果,感觉确实还不错。基于此,我重构了源代码,主要工作:剔除了源代码对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:
PyTorch模型导出ONNX中间格式的方法,熟悉可视化ONNX节点工具: netron的使用;
推理引擎 ONNX Runtime 和 TensorRT的安装及使用方法,包括:TensorRT python API 和 C++ API使用;
熟悉TensorRT 工具Polygraphy的python API使用方法和基本的命令行调用指令;
基本的CUDA编程知识:核函数的编写与启动,常用的内存模型:全局内存、共享内存,etc;
Makefile编程语法,C++编译规则和nvcc 编译cuda程序规则;
通过这篇文章你将了解并学习到以下进阶知识:
模型转换完,如何确保数据流的无损传递以及推理结果的一致性验证方法,即:如何debug问题和校验结果准确性 (我想这是大家最关心的部分,毕竟网上大几百的课程也未必会详细的告诉大家这些方法);
如何在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 ====================================================================================================================
部署的大体思路如下:
Sparse4Dv3模型部署的过程需要将我们将训练好的模型
.pth 文件,转换为中间文件
.onnx,最后转换为*.engine 文件。该过程需要解决PyTorch模型与ONNX框架算子的兼容性以及模型运行加速两大需求; 模型pth下载链接:https://drive.google.com/file/d/1sSMNB7T7LPKSr8nD9S_tSiu1mJrFMZ1I/view?usp=sharing**
Sparse4Dv3 PyTorch模型实际上就是一个计算图。模型部署时通常需要我们将模型转换成静态的计算图,即没有控制流(分支语句、循环语句)的计算图,这点很重要,Sparse4Dv3 Head模型转换过程我们就会遇到这个问题;
PyTorch 框架自带对 ONNX 的支持,只需要我们构造一组随机的输入,并对模型调用 torch.onnx.export 即可,完成 PyTorch 到 ONNX 的转换。
*
推理引擎 ONNX Runtime 对 ONNX 模型有原生的支持,提供了python API 和 C++ API。给定一个 .
onnx 文件,只需简单使用 ONNX Runtime 的 Python API 就可以完成模型推理。依据上述工具的使用,我们可以完成PyTorch到ONNX的推理一致性验证。
推理引擎 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算子将原始的3D关键点:固定关键点和可学习关键点,与图像特征对齐,从而提取有效的图像特征
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)
针对上述问题,作者提出了加速的需求,并给出了加速方案:
图四:使用CUDA多线程并行化加速DFA特征聚合
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算子怎么部署呢?具体开发流程是什么?
这里先说结论:
首先,在PyTorch代码中使用符号函数:symbolic:构建DFA自定义算子PyTorch 到 ONNX 映射规则;
其次,将带有自定义算子DFA,PyTorch 到 ONNX 映射规则的*.pth模型文件转化为ONNX中间文件;
最后,基于指定版本的TensorRT C++ 注册plugin的 API,构建自定义DFA plugin,构建编译规则,编译该plugin,在本地将会生成动态库*.so文件;
对比验证PyTorch推理结果和TensorRT推理结果,确保结果的一致性;
通过trtexec command line 加载上述动态库到目标trt engine中,生成带有plugin的目标模型engine;
这里,不得不提下我之前幻想能够实现部署目标的几个解决方案的思路,以及后续为什么我最终敲定了上述方案:
被我舍弃的思路一:
拆分法
由于DFA是自定义算子,这意味着:PyTorch的ATen[5]基本库没有官方实现、ONNX节点Op不支持该算子、TensorRT(包含官方实现的Plugin)不支持该算子。因此,在转换ONNX过程必定面临失败,更别提转换TensorRT Engine了。那么比较容易想到的方法是:在自定义算子 出现的位置,我们把模型逐个拆分开,最后,在C++车端代码中逐个把模型串联起来不就可以了?一般来说这种方案没啥问题,但在Sparse4Dv3中,我们不推荐这么做,请看下面:
图五:Sparse4D Head 单次推理DFA出现次数
很显然,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节点映射关系:
图六:使用symbolic符号函数以及g.op()搭建算子映射关系
第二步:生成ONNX中间文件。有了映射关系,我们就可以生成ONNX文件了,netron可视化如下:
图七:Custom op:DeformableAttentionAggrPlugin ONNX node
这里的ONNX文件是无法成功转换trt engine,因为,ONNX中的DeformableAttentionAggrPlugin节点只是一个接口,没有对应的TensorRT Op实现的。因此,下一步需要在TensorRT中注册该算子,并将核心实现填充进去。
第三步:Custom Operator Plugin: DeformableAttentionAggrPlugin注册(我用最简单的话来解释):
打开浏览器搜索github 仓库 TensorRT,指定版本8.5,打开plugin文件夹,这里以官方的实现的第一个plugin: batchTilePlugin为例:
图八:以batchTilePlugin为例模仿搭建DeformableAttentionAggrPlugin
我们需要构建四个文件分别为:
deformableAttentionAggrPlugin.h
deformableAttentionAggrPlugin.cpp
deformableAttentionAggrPlugin.cu
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类类常用方法:
图九:ICudaEngine类的实例对象就是我们反序列化后生成的engine对象
图十:engine.get_binding_dtype获得绑定的tensor类型,它和ONNX输入输出类型保持一致
图十一:engine.get_binding_name获得绑定的tensor名称,它和ONNX输入输出名称保持一致
I EaecutionConteat类常用方法:
图十二:IExecutionContext实例化:context = engine.create_execution_context()
图十三:get_bind_shape 获得execute_v2接口完成模型的推理
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))} ."