pytorch 28 jit模型（TorchScript）转换为onnx模型

万里鹏程转瞬至

已于 2022-06-06 11:42:59 修改

阅读量2.5k

点赞数 1

分类专栏： pytorch工程实践文章标签： pytorch 深度学习神经网络

于 2022-06-06 11:05:28 首次发布

本文链接：https://blog.csdn.net/a486259/article/details/125141357

版权

pytorch工程实践专栏收录该内容

51 篇文章 ¥59.90 ¥99.00

订阅专栏

本文介绍了如何将PyTorch的TorchScript（jit）模型转换为ONNX模型，以适应不同的部署需求。通过创建jit模型、转化为ScriptModule、设置example_outputs参数来完成转换，并使用onnxruntime进行测试，验证了转换前后模型输出的一致性。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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,

了解本专栏