pytorch 28 jit模型(TorchScript)转换为onnx模型

TorchScript是pytorch模型c++下libtorch部署的必须要求,但是已经部署好的jit模型如果有一天突然想换成onnx部署、openvion部署,tensorRT部署则存在一定的困难。因为,pytorch是不支持直接将jit模型模型转换为onnx模型的。为此,对jit模型转换为onnx模型进行研究实现。这里以多输入多输出的jit模型为例,实现将jit模型转换为onnx模型,并用onnxruntime进行测试。测试结果表明,在相同的输入数据下,jit模型与onnx模型的输出结果是完全一模一样的。

1、创建jit模型

如果有现成的jit模型不需要这一步。这里是构建多输入多输出模型,并保存为jit模型

import torch
import torch.nn as nn
import torch.nn.functional as F
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(6, 16, kernel_size=1, stride=1, padding=0)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(16, 16, kernel_size=1, stride=1, padding=0)
        self.conv31 = nn.Conv2d(16, 3, kernel_size=1, stride=1, padding=0)
        self.conv32 = nn.Conv2d(16, 3, kernel_size=1, stride=1, padding=0)
   

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
上一篇 2022年6月8日 上午10:53
下一篇 2022年6月8日 上午10:56

相关推荐

本站注重文章个人版权,不会主动收集付费或者带有商业版权的文章,如果出现侵权情况只可能是作者后期更改了版权声明,如果出现这种情况请主动联系我们,我们看到会在第一时间删除!本站专注于人工智能高质量优质文章收集,方便各位学者快速找到学习资源,本站收集的文章都会附上文章出处,如果不愿意分享到本平台,我们会第一时间删除!