使用pytorch输出onnx模型并转换成openvino的IR模型

最新推荐文章于 2024-09-11 17:04:06 发布

ZHIKELIN

最新推荐文章于 2024-09-11 17:04:06 发布

阅读量3.5k

点赞数

分类专栏： onnx pytorch IR 文章标签： pytorch python

本文链接：https://blog.csdn.net/ZHIKELIN/article/details/107314475

版权

本文介绍了如何在Windows 10环境下，使用PyTorch 1.5.1和torchvision 0.6.1构建的模型转换为ONNX格式，并进一步转换为OpenVINO的IR模型。关键步骤包括解决因包版本不一致导致的错误，特别是将networkx降级到2.0以避免转换过程中出现的问题。最终通过`mo_onnx.py`脚本完成IR模型的生成。

摘要由CSDN通过智能技术生成

经过好几天的试错，终于成功转换出模型了，出错主要问题是各种包版本不一致的问题。

安装pytorch/torchvision

安装环境和包版本：
Windows 10
torch：1.5.1；
torchvision：0.6.1；
python：3.7.6
onnx：1.7.0
openVino：Anaconda3-2020.02-Windows-x86_64.exe
tensorflow>=1.2.0,<2.0.0
mxnet>=1.0.0,<=1.3.1
networkx=2.0
numpy=1.19.0
protobuf3.6.1
test-generator0.1.1
defusedxml>=0.5.0

anaconda3环境安装：
conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
安装openVino时，默认会按照如下版本安装
tensorflow>=1.2.0,<2.0.0
mxnet>=1.0.0,<=1.3.1
networkx>=1.11
numpy>=1.12.0
protobuf3.6.1
onnx>=1.1.2
test-generator0.1.1
defusedxml>=0.5.0
所以，如果没有设置最高版本都会按最新版本安装，但是会导致版本不一致，其中networkx的版本默认安装是2.4，但生成IR模型时出现如下错误：

[ ERROR ]  Cannot pre-process ONNX graph after reading from model file 
"D:\project\Python\detection\SSD-Tensorflow0713\super_resolution.onnx". 
File is corrupt or has unsupported format. Details: 'Graph' object has no 
attribute 'node'. For more information please refer to Model Optimizer FAQ 
 (https://docs.openvinotoolkit.org/latest/_docs_MO_DG_prepare_model_Model_Optimizer_FAQ.html), question #44.

将networkx的版本改成2.0就适配了，我使用的是pip安装：

pip install networkx==2.0

使用pytorch输出onnx模型并转换成openvino的IR模型

model–>onnx–>IR转换

import io
import numpy as np

from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
import torch.nn as nn
import torch.nn.init as init
import onnx
import onnxruntime


class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

        self._initialize_weights()

    def forward(self, x):
        x = self.relu