tensorRT的插件(自定义算子)实现

插件实现

算子就是指封装好的一个计算函数。算子的基本组成就是一些加减乘除和一些基本运算,封装成算子的话,可以方便后续的重复调用,软工上的思想就是抽象和复用。粒度太细完全无法沟通。AI compiler就是某种意义上的直接展开用。自定义算子(插件)就是自己把一些常用的计算过程封装成一个算子,这个封装官方并没有实现过。比如nn.Conv2d和nn.ReLU都是官方的算子。当然你也可以选择不做成算子,按照数学公式实现出来进行部署同样不会报错,但就是导出的onnx非常乱。还有的情况就是forward中包含了自己用cpp写的cuda实现时,这是不能被trace的,即PyTorch 无法 trace (跟踪)数据流动,这样就导不出计算图,所以这时也需要把这一坨弄成一个plugin,就能够trace了。或者全换成pytorch实现。

自定义算子
PyTorch转ONNX踩坑记

插件实现要注意以下几点

  1. 如何在pytorch里面导出一个插件。因为模型是通过pytorch导onnx导出来的,因此我们需要在pytorch中描述下想导出的插件是什么类型,输入是什么东西,有什么参数。
  2. 插件解析时如何对应,在onnx parser中如何处理。即如何解析这个onnx让tensorRT知道这个插件。
  3. 插件的creator实现,这个creator是tensorRT里面的一个概念,负责创建插件的具体实现这个类
  4. 插件的具体实现,继承自IPluginV2DynamicExt
  5. 插件的序列化与反序列化,因为插件在编译的时候可以去使用这个插件,在编译过后会确定好插件选择什么参数(即tensorRT在编译过程中取选择最优配置),然后这些参数要被序列化储存成文件,下次使用时又要反序列化去把文件加载成参数。

你可能感兴趣的:(模型部署,深度学习,人工智能,部署)