【模型加速】自定义TensorRT NMS3D插件(1)

需求是这样的，在做PointPillars模型的加速的时候我注意到网络的检测头部分小型操作很多，加速效果不明显。此外，3D检测模型的NMS部分通常是作为后处理的一部分来单独实现，TensorRT并没有直接支持3D NMS的导出。本着学习的目的，我将PointPillars模型中的检测头(单头)和3D NMS两部分合并到一个TensorRT Plugin，实现端到端的推理。其最终效果如下右图所示，自定义的NMS3D Plugin包含了整个后处理部分。

如何在onnx的输出后面增加NMS3D节点?

这一步涉及到修改onnx模型，可借助TensorRT自带的小工具ONNX GraphSurgeon来完成。它可以增加或者移除某些onnx节点，修改名字或者维度等等。ONNX GraphSurgeon工具的安装也很简单，先安装nvidia-pyindex，然后再安装onnx-graphsurgeon。

pip install nvidia-pyindex
pip install onnx-graphsurgeon

然后再是修改计算图的操作，我这里给出两种实现方式仅供参考。

方法一：

# Here we'll register a function to do all the subgraph-replacement heavy-lifting.
# NOTE: Since registered functions are entirely reusable, it may be a good idea to
# refactor them into a separate module so you can use them across all your models.
@gs.Graph.register()
def add_nms3d(self, inputs, outputs):# Disconnect output nodes of all input tensorsfor inp in inputs:inp.outputs.clear()### Disconnet input nodes of all output tensorsfor out in outputs:out.inputs.clear()attrs = collections.OrderedDict()attrs['anchor_sizes'] = anchor_sizesattrs['anchor_bottom_heights'] = anchor_bottom_heights# Insert the new node.return self.layer(op="NMS3D", inputs=inputs, outputs=outputs, name="nms3d", attrs=attrs)def simplify_onnx():model = onnx.load("pointpillar_raw.onnx")graph = gs.import_onnx(onnx_model)tmap = graph.tensors()inputs = [tmap['cls_preds'],tmap['box_preds'],tmap['dir_cls_preds']]outputs = [gs.Variable(name="nms3d_output", dtype=np.float32, shape=(1,100,9))]graph.add_nms3d(inputs, outputs)graph.outputs = outputsgraph.cleanup()graph.toposort()onnx.save(model_simplify, "pointpillar_simplify.onnx")print("export ok...")

方法二:

def simplify_onnx():#model = onnx.load("pointpillar_raw.onnx")model = onnx.load("pointpillar_fcn_max_nchw_cudapp.onnx")while len(model.graph.output):model.graph.output.remove(model.graph.output[0])model.graph.output.extend([onnx.helper.make_tensor_value_info('nms3d_output', onnx.TensorProto.FLOAT, [1,100,9]),])  attrs = collections.OrderedDict()attrs['anchor_sizes'] = anchor_sizesattrs['anchor_bottom_heights'] = anchor_bottom_heightsgraph = gs.import_onnx(model)tmap = graph.tensors()inputs = [tmap['cls_preds'],tmap['box_preds'],tmap['dir_cls_preds']]outputs = [tmap['nms3d_output']]nms3d_layer = graph.layer(op="NMS3D", inputs=inputs, outputs=outputs, name="nms3d", attrs=attrs)graph.cleanup()graph.toposort()onnx_module = gs.export_onnx(graph)onnx.save(onnx_module, "pointpillar_simplify.onnx")print("export ok...")

【参考文献】

TensorRT详细入门指北，如果你还不了解TensorRT，过来看看吧！ - 知乎

安装onnx-graphsurgeon_人类高质量算法工程师的博客-CSDN博客

如何修改已有的ONNX模型 - 知乎

Polygraphy逐层对比onnx和tensorrt模型的输出 - 知乎

深度学习系列4：onnx_IE06的博客-CSDN博客_onnx学习

【模型加速】自定义TensorRT NMS3D插件(1)相关推荐

【模型加速】TensorRT详解
■ TensorRT概述 NVIDIA®TensorRT™的核心是一个C++库,可以促进在NVIDIA图形处理单元(GPU)上的高性能推断.它旨在与Tensorflow.Caffe.Pytorch.M ...
【模型加速】TensorRT安装、测试及常见问题
■ 安装过程一.安装依赖环境 ● Ubuntu 20.04 ● CUDA 11.1 ● cuDNN 8.0.4 ● python 3.8.5 – 可以通过命令查看cuda.cudnn.python版 ...
yolo模型部署——tensorRT模型加速+triton服务器模型部署
将最近的工作做个记录,方便日后学习回顾: 1.针对项目需求开发满足任务的模型,拿到任务就要去选相应的算法,由于是工程应用型,必须找填坑多的算法,这样遇到问题可参考的资料多. 2.做好以后,还要将开发的 ...
Ultralytics公司YOLOv8来了(训练自己的数据集并基于NVIDIA TensorRT和华为昇腾端到端模型加速)--跟不上“卷“的节奏
Official YOLOv8 训练自己的数据集并基于NVIDIA TensorRT和华为昇腾端到端模型加速说明: 本项目支持YOLOv8的对应的package的版本是:ultralytics-8. ...
php多选筛选,DEDECMS自定义模型筛选多选版插件
DEDECMS自定义模型筛选多选版插件,像分类信息网站一样的筛选功能. 一.文件夹说明: incluede 核心函数目录二.安装说明 1.把这些文件夹全部复制到根目录粘贴,或者按文件 ...
win10下 yolov8 tensorrt模型加速部署【实战】
Windows10下yolov8 tensorrt模型加速部署[实战] TensorRT-Alpha基于tensorrt+cuda c++实现模型end2end的gpu加速,支持win10.linux ...
【模型加速】PointPillars模型TensorRT加速实验(7)
按照[模型加速]PointPillars模型TensorRT加速实验(7)中给出的思路对已有的推理代码进行优化,简而言之就是保持数据在GPU显存中流动,尽量避免内存和显存之间的流动. PFN推理v2 ...
模型加速之INT8量化原理及实践（基于TensorRT）
一.模型量化: 1.量化的定义是将网络参数从Float-32量化到更低位数,如Float-16.INT8.1bit等. 2.量化的作用:更小的模型尺寸.更低的功耗.更快的计算速度.下图是不同数据结构比 ...
tensorrt轻松部署高性能dnn推理_部署环境之：tensorRT的插件
TensorRT是一个高性能的深度学习推理(Inference)优化器,可以为深度学习应用提供低延迟.高吞吐率的部署推理.TensorRT可用于对超大规模数据中心.嵌入式平台或自动驾驶平台进行推理加速 ...

【模型加速】自定义TensorRT NMS3D插件(1)

如何在onnx的输出后面增加NMS3D节点?

【模型加速】自定义TensorRT NMS3D插件(1)相关推荐

最新文章

热门文章