tensorRT 自定义算子plugin的实现_tensorrt 支持自定义plugin

作者：菜鸟追梦旅行 | 2024-06-17 05:57:11

踩

tensorrt 支持自定义plugin

文章目录

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

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/菜鸟追梦旅行/article/detail/729885