模型部署——自定TensorRT算子

CenterPoint基于OpenPcDet已经导出一个完整的ONNX，转onnx可以参考：模型部署——CenterPoint转ONNX(自定义onnx算子)，下面我们基于导出的onnx模型，并用TensorRT进行推理，直接推理有些算子不支持，下面讲解下TensorRT自定义算子，由于torch.unique相比，其他的几个算子相对比较简单点，我会重点讲解下自定义Unique的TensorRT算子，torch.unique功能如下：

        # sorted:是否返回无重复张量按照数值进行排序，默认是升序排列，sorted并非表示降序
        # return_inverse:是否返回原始张量中每个元素在处理后的无重复张量中对应的索引
        # return_counts：统计原始张量中每个独立元素的个数
        # dim:值沿那个维度进行unique的处理
        unq_coords, unq_inv, unq_cnt = torch.unique(merge_coords, return_inverse=True, return_counts=True, dim=0)

注意：该算子没有通用性，只是实现上述python代码的功能，与torch.unique的三个输出相对应，考虑通用性，需要进一步实现各个参数选项对应的功能。

核心难点在于思路，有思路后顺利用cuda实现就行，所以先介绍Unique算子的实现思路如下：

• 排序：先对网格坐标coords进行排序，cuda实现排序可以参考：CUDA动态并行实现快速排序，也可以直接调用CUDA Thrust库。记排序后的坐标为coord_sort
• 去重复：对排序后的坐标coord_sort去除重复元素，可以自己写，也可以直接调用Thrust unique函数，得到去重后的坐标unq_coords
• 哈希：unq_coords的索引下标和对应的值互换，用一个哈希数组进行存储，就可以根据去重后的网格坐标，获取每个元素在处理后的无重复张量中对应的索引，
• 遍历：遍merge_coords，通过哈希数组可以得到原始张量中每个元素在处理后的无重复张量中对应的索引

接下来将要讲解的TensorRT，将会是基于最新的8.4版本。主要是讲解：

• TensorRT自定义插件的使用方式
• 如何添加自己的自定义算子

有兴趣的可以先看看TensorRT的官方文档，官方文档的介绍简单意骇

首先按照官方plugin的排布方式，下面随便挑了个官方plugin，在官方提供基础算子修改，会节省很多时间