tensorRT 自定义算子plugin的实现

作者：AllinToyou | 2024-06-13 02:54:53

踩

文章目录

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

1
2
3
4
5
6
7
8
9

本文内容由网友自发贡献，转载请注明出处：https://www.wpsshop.cn/w/AllinToyou/article/detail/710622