文章目录
在实际部署模型时,由于现有的算子无法满足需求,经常需要
自定义
算子并导出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