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. 完整代码下载
在实际部署模型时,由于现有的算子无法满足需求，经常需要
自定义算子并导出onnx文件，并在tensort中通过实现plugin来支持自定义算子的推理。下面给出个具体应用中的示例：
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