tensorRT 自定义算子plugin的实现

文章目录

• • 1. 自定义算子导出onnx
  • • 1.1 自定义算子的实现类
    • 1.2 自定义算子类
    • 1.3 onnx导出
  • 2. 自定义算子plugin的实现
  • • 2.1 Plugin类的实现
    • • 2.1.1 构造函数的实现
      • 2.1.2 析构函数实现
      • 2.1.3 serialize函数的实现
      • 2.1.4 destory函数的实现
      • 2.1.5 enqueue函数的实现（关键）
      • • 2.1.5.1 customScalarImpl核函数的实现
      • 2.1.6 clone函数的实现
      • 2.1.7 supportsFormatCombination函数的实现
      • 2.1.8 get相关的函数
      • 2.1.9 其他函数
      • 总结
    • 2.2 PluginCreator 类的实现
    • • 2.2.1 构造函数实现
      • 2.2.2 析构函数实现
      • 2.2.3 createPlugin函数实现
      • 2.2.4 deserializePlugin 函数实现
      • 2.2.5 get相关函数的实现
      • 2.2.6 set相关函数的实现
      • 总结
  • 3. 自定义plugin案例
  • • 3.1 Plugin类
    • • 3.1.1 类的定义
      • 3.1.2 类的实现
    • 3.2 PluginCreator类
  • 3. 完整代码下载

1. 自定义算子导出onnx

实现一个自定义算子CustomScalar, 对输入的x， 先加一个数r,再乘以s，得到对应的输出：(x + r) * s

import torch
import torch.onnx
import torch.nn as nn
import onnxruntime
import onnx
import onnxsim
import os
from collections import OrderedDict