pytorch自定义算子并导出onnx计算图详细代码教程

最新推荐文章于 2024-07-04 22:50:27 发布
最新推荐文章于 2024-07-04 22:50:27 发布
阅读量1.5k 收藏 4
点赞数
文章标签： pytorch 人工智能 python
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/m0_37870649/article/details/132702392
版权 
(torch_env) root@f123:/mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin# python gen-onnx.py
对应opset文件夹代码在这里： /root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx
inference output =
tensor([[[[3.9281, 5.9852, 3.9281],
          [5.9852, 8.9989, 5.9852],
          [3.9281, 5.9852, 3.9281]]],


        [[[0.7311, 2.8577, 2.8577],
          [2.8577, 3.9281, 2.8577],
          [2.8577, 2.8577, 0.7311]]]], grad_fn=<MYSELUImplBackward>)
/root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx/utils.py:90: UserWarning: 'enable_onnx_checker' is deprecated and ignored. It will be removed in the next PyTorch release. To proceed despite ONNX checker failures, catch torch.onnx.ONNXCheckerError.
  warnings.warn("'enable_onnx_checker' is deprecated and ignored. It will be removed in "
==================================call symbolic
Warning: Unsupported operator Plugin. No schema registered for this operator.
Warning: Unsupported operator Plugin. No schema registered for this operator.
Warning: Unsupported operator Plugin. No schema registered for this operator.
graph(%image : Float(*, 1, *, *, strides=[9, 9, 3, 1], requires_grad=0, device=cpu),
      %conv.weight : Float(1, 1, 3, 3, strides=[9, 9, 3, 1], requires_grad=1, device=cpu),
      %conv.bias : Float(1, strides=[1], requires_grad=1, device=cpu),
      %myselu.param : Float(3, strides=[1], requires_grad=1, device=cpu)):
  %4 : Float(*, 1, *, *, strides=[9, 9, 3, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%image, %conv.weight, %conv.bias) # /root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/nn/modules/conv.py:442:0
  %5 : Float(3, strides=[1], device=cpu) = onnx::Constant[value= 3  2  1 [ CPUFloatType{3} ]]()
  %output : Float(1, 1, 3, 3, strides=[9, 9, 3, 1], requires_grad=1, device=cpu) = onnx::Plugin[info="{\"attr1_s\": \"\0-0\/0/\/,)\0-.\/+0\/./\0--\/.-\/+/\0-/\/.,\/-.\0-,\//0\//*\0--\//)\/,.\0-.\.00\/-/\", \"attr2_i\": [1, 2, 3], \"attr3_f\": 222}", name="MYSELU"](%4, %myselu.param, %5) # /mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin/gen-onnx.py:38:0
  return (%output)

Traceback (most recent call last):
  File "/root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx/utils.py", line 785, in _export
    _check_onnx_proto(proto)
RuntimeError: No Op registered for Plugin with domain_version of 11

==> Context: Bad node spec for node. Name: Plugin_2 OpType: Plugin

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin/gen-onnx.py", line 80, in <module>
    torch.onnx.export(
  File "/root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx/__init__.py", line 316, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx/utils.py", line 107, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx/utils.py", line 787, in _export
    raise ONNXCheckerError(e)
torch.onnx.utils.ONNXCheckerError: No Op registered for Plugin with domain_version of 11

==> Context: Bad node spec for node. Name: Plugin_2 OpType: Plugin
(torch_env) root@f123:/mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin#
(torch_env) root@f123:/mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin#
(torch_env) root@f123:/mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin#
(torch_env) root@f123:/mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin# python gen-onnx.py
对应opset文件夹代码在这里： /root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx
inference output =
tensor([[[[3.9281, 5.9852, 3.9281],
          [5.9852, 8.9989, 5.9852],
          [3.9281, 5.9852, 3.9281]]],


        [[[0.7311, 2.8577, 2.8577],
          [2.8577, 3.9281, 2.8577],
          [2.8577, 2.8577, 0.7311]]]], grad_fn=<MYSELUImplBackward>)
/root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx/utils.py:267: UserWarning: `add_node_names' can be set to True only when 'operator_export_type' is `ONNX`. Since 'operator_export_type' is not set to 'ONNX', `add_node_names` argument will be ignored.
  warnings.warn("`{}' can be set to True only when 'operator_export_type' is "
/root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/onnx/utils.py:267: UserWarning: `do_constant_folding' can be set to True only when 'operator_export_type' is `ONNX`. Since 'operator_export_type' is not set to 'ONNX', `do_constant_folding` argument will be ignored.
  warnings.warn("`{}' can be set to True only when 'operator_export_type' is "
==================================call symbolic
Warning: Unsupported operator Plugin. No schema registered for this operator.
Warning: Unsupported operator Plugin. No schema registered for this operator.
Warning: Unsupported operator Plugin. No schema registered for this operator.
graph(%image : Float(*, 1, *, *, strides=[9, 9, 3, 1], requires_grad=0, device=cpu),
      %conv.weight : Float(1, 1, 3, 3, strides=[9, 9, 3, 1], requires_grad=1, device=cpu),
      %conv.bias : Float(1, strides=[1], requires_grad=1, device=cpu),
      %myselu.param : Float(3, strides=[1], requires_grad=1, device=cpu)):
  %4 : Float(*, 1, *, *, strides=[9, 9, 3, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%image, %conv.weight, %conv.bias) # /root/anaconda3/envs/torch_env/lib/python3.9/site-packages/torch/nn/modules/conv.py:442:0
  %5 : Float(3, strides=[1], device=cpu) = onnx::Constant[value= 3  2  1 [ CPUFloatType{3} ]]()
  %output : Float(1, 1, 3, 3, strides=[9, 9, 3, 1], requires_grad=1, device=cpu) = onnx::Plugin[info="{\"attr1_s\": \"\0-0\/0/\/,)\0-.\/+0\/./\0--\/.-\/+/\0-/\/.,\/-.\0-,\//0\//*\0--\//)\/,.\0-.\.00\/-/\", \"attr2_i\": [1, 2, 3], \"attr3_f\": 222}", name="MYSELU"](%4, %myselu.param, %5) # /mnt/d/work/learning-cuda-trt/learning-cuda-trt-main/tensorrt-basic-1.8-integrate-easyplugin/gen-onnx.py:38:0
  return (%output)

Done.!
(torch_env)

解决：

    # enable_onnx_checker=False

更改为：

    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK

pytorch自定义算子并导出onnx计算图详细代码教程_operatorexporttypes_蛇皮小娃娃的博客-CSDN博客

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import torch.autograd
import os

class MYSELUImpl(torch.autograd.Function):

    # reference: https://pytorch.org/docs/1.10/onnx.html#torch-autograd-functions
    @staticmethod
    def symbolic(g, x, p):
        print("==================================call symbolic")
        return g.op("MYSELU", x, p, 
            g.op("Constant", value_t=torch.tensor([3, 2, 1], dtype=torch.float32)),
            attr1_s="这是字符串属性", 
            attr2_i=[1, 2, 3], 
            attr3_f=222
        )

    @staticmethod
    def forward(ctx, x, p):
        return x * 1 / (1 + torch.exp(-x))


class MYSELU(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.param = nn.parameter.Parameter(torch.arange(n).float())

    def forward(self, x):
        return MYSELUImpl.apply(x, self.param)


class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv = nn.Conv2d(1, 1, 3, padding=1)
        self.myselu = MYSELU(3)
        self.conv.weight.data.fill_(1)
        self.conv.bias.data.fill_(0)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.myselu(x)
        return x


# 这个包对应opset11的导出代码，如果想修改导出的细节，可以在这里修改代码
# import torch.onnx.symbolic_opset11
print("对应opset文件夹代码在这里：", os.path.dirname(torch.onnx.__file__))

model = Model().eval()
input = torch.tensor([
    # batch 0
    [
        [1,   1,   1],
        [1,   1,   1],
        [1,   1,   1],
    ],
        # batch 1
    [
        [-1,   1,   1],
        [1,   0,   1],
        [1,   1,   -1]
    ]
], dtype=torch.float32).view(2, 1, 3, 3)

output = model(input)
print(f"inference output = \n{output}")

dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(
    model, 

    # 这里的args，是指输入给model的参数，需要传递tuple，因此用括号
    (dummy,), 

    # 储存的文件路径
    "workspace/demo.onnx", 

    # 打印详细信息
    verbose=True, 

    # 为输入和输出节点指定名称，方便后面查看或者操作
    input_names=["image"], 
    output_names=["output"], 

    # 这里的opset，指，各类算子以何种方式导出，对应于symbolic_opset11
    opset_version=10, 

    # 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1
    # 通常，我们只设置batch为动态，其他的避免动态
    dynamic_axes={
        "image": {0: "batch", 2: "height", 3: "width"},
        "output": {0: "batch", 2: "height", 3: "width"},
    },

    # 对于插件，需要禁用onnx检查
    # enable_onnx_checker=False
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK

)

print("Done.!")

 

a flying bird
关注
复杂算子onnx导出(1): trace的实现
@bangbang的博客
04-23 210
通过lidar传感器采集的点云数据，它是非常稀疏的, 在分布的空间中只占1% ~5%, 如果使用普通的3D卷积去计算，大部分都是在卷空气，因此提出了使用3D 卷积的方式，只在tensor中有值的位置进行卷积计算，大大提高了卷积的效率，同时将输入输出进行稀疏化存储。稀疏卷积存在两种模式:(1)标准的稀疏卷积（2）subm模式的稀疏卷积标准稀疏卷积：它的计算和普通的卷积计算是一样的，只不过它将输入和输出进行稀疏化存储。通过features。
Pytorch图像分类实战笔记Task01-Task07|Datawhale组队学习
weixin_45480097的博客
01-17 859
将数据集统一命名为数字-->获得文件名列表full_dataset-->打乱列表random_dataset-->set train_rate = 0.8,train_dataset = random_dataset[:int(full_dataset)*0.8],test_dataset =random_dataset[int(full_dataset)*0.8:] -->移动文件。文件：gif格式的、非三通道图片（三通道图可以是彩色图，可以是灰度模式的图像。第一步，导入所需要的库。
自定义算子onnx-bin
04-01
自定义算子onnx-bin
ONNX，从了解到精通
qq_37794738的博客
04-01 1738
这篇教程用到的技术很多，需要有一定的基础，博主建议先学会添加算子的方法，走通整个流程之后，再回来填那些坑，比如如何添加Pytorch的具体实现插件，有没有更便捷的添加ONNX插件的方法等等，这些后面博主都会更新；
【pytorch扩展】CUDA自定义pytorch算子（简单demo入手）
最新发布
木盏
07-04 410
Pytorch作为一款优秀的AI开发平台，提供了完备的自定义算子的规范。我们用torch开发时，经常会因为现有算子的不足限制我们idea的迸发。众所周知，CUDA/C++都是编译性语言，编译以后再调用会比python这种解释性语言更快。所以，我们需要对CUDA有一个编译过程。直接“python setup.py install”即可完成cuda算子的编译和安装。这里的"sum_double_op"就是我们用CUDA写的算子。那这个可以直接调用，并且可以传递梯度的算子，需要怎么做呢？
pytorch自定义算子 native_functions.yaml
weixin_43561808的博客
08-17 2365
pytorch自定义算子 native_functions.yaml 在pytorch的文件夹中搜索native_functions.yaml，可以看到，所有pytorch原生的函数都是在这里注册的 在同文件夹下，存在一个README.md，这是官方教程，教我们怎么用native_functions.yaml注册自己自定义的算子 下面我仿照pytorch原生的relu算子，写一个自己的myrelu 首先，找到native_functions.yaml，有几个函数是我们可以参考的 - func: relu(T
记录pytorch实现自定义算子并转onnx文件输出
Briwisdom的博客
11-14 1218
概览：记录了如何自定义一个算子，实现pytorch注册，通过C++编译为库文件供python端调用，并转为onnx文件输出
Pytorch 自定义算子
绿色羽毛
08-13 1917
自定义激活函数算子 Pytorch 自定义激活函数，都需要继承torch.autograd.Function 类，其内部需要定义两个静态方法【@staticmethod】：forward & 1.forward方法：定义了输入张量（可以是多个），返回对应的输出（也可以是多个张量）；2. backward方法：定义了输出的剃度张量（可以是多个，必须和输出张量一一对应），返回输入张量（可以是多个，必须喝输入张量一一对应）【之所以一一对应，是因为计算图中的每个张量在方向传播的时候，输出张量和输入张量对应
PyTorch扩展自定义算子
weixin_45467930的博客
06-29 1041
pytorch拓展自定义算子方法总结
Pytorch模型部署：torch-onnx-ncnn
MY_WARNING的博客
08-22 2108
部署
PyTorch与ONNX模型转换实践指南
![PyTorch与ONNX模型转换实践指南](https://img-blog.csdnimg.cn/img_convert/1614e96aad3702a60c8b11c041e003f9.png) ...PyTorch与ONNX模型转换使开发人员能够将PyTorch模型导出为ONNX格式，从而实现跨框架部
pytorch转onnx自定义算子
weixin_43656490的博客
12-21 1009
pytorch导出含有自定义算子的onnx文件
ONNX自定义算子与文件导出流程再梳理
Briwisdom的博客
11-07 1365
工作需要把onnx导出文件的流程又梳理了一遍，弄清楚了如何自定义算子，以及导出结果检查。分享一下~
PyTorch自定义CUDA算子教程与运行时间分析
qq_35985044的博客
08-23 916
最近因为工作需要，学习了一波CUDA。这里简单记录一下PyTorch自定义CUDA算子的方法，写了一个非常简单的example，再介绍一下正确的PyTorch中CUDA运行时间分析方法。 所有的代码都放在了github上，地址是：https://github.com/godweiyang/torch-cuda-example 完整流程 下面我们就来详细了解一下PyTorch是如何调用自定义的CUDA算子的。 首先我们可以看到有四个代码文件： main.py，这是python入口，也就是你平时写
PyTorch扩展自定义PyThon/C++(CUDA)算子的若干方法总结
小白学视觉
10-21 839
点击上方“小白学视觉”，选择加"星标"或“置顶”重磅干货，第一时间送达作者丨奔腾的黑猫@知乎来源丨https://zhuanlan.zhihu.com/p/158643792本文仅用于学术分享，如有侵权，请联系后台作删文处理。导读关于PyTorch构建扩展的一些基础操作，官方往往已经出具了完整的教程。本文对这些官方教程的链接进行了整理，以供读者查阅。在做毕设的时候，需要实现一个PyTorch原生...
导出onnx模型(1)：torch.onnx.export 多输入
weixin_43470772的博客
05-06 718
将模型的多个输入构造为dict，输入到torch.onnx.export（）函数中，实现多输入模型导出onnx。
评论 1
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值