点击上方蓝字关注我们

微信公众号：OpenCV学堂

关注获取更多计算机视觉与深度学习知识

Pytorch ONNX格式支持

ONNX是一种深度学习权重模型的表示格式，ONNX格式可以让AI开发者在不同框架之间相互转换模型，实现调用上的通用性。当前PyTorch*, Caffe2*, Apache MXNet*, Microsoft Cognitive Toolkit* 、百度飞桨都支持ONNX格式。OpenVINO的模型优化器支持把ONNX格式的模型转换IR中间层文件。

当前OpenVINO官方支持的ONNX模型主要包括：bert_large，bvlc_alexnet，bvlc_googlenet，bvlc_reference_caffenet，bvlc_reference_rcnn_ilsvrc13 model，inception_v1，inception_v2，resnet50，squeezenet，densenet121，emotion_ferplus，mnist，shufflenet，VGG19，zfnet512。需要注意的是这些模型升级版本并不被支持。

从OpenVINO的2019R04版本开始支持所有公开的Pytorch模型，支持的模型列表如下：

Pytorch ONNX到OpenVINO IR转换

下面的例子演示了如何从torchvision的公开模型中转换为ONNX，然后再转换为IR，使用OpenVINO完成调用的完整过程。我们将以resnet18为例来演示。

01

下载模型与转ONNX格式

要下载与使用torchvision的预训练模型，首选需要安装好pytorch，然后执行下面的代码就可以下载相关支持模型：

1import torchvision.models as models  2resnet18 = models.resnet18(pretrained=True)  3alexnet = models.alexnet(pretrained=True)  4squeezenet = models.squeezenet1_0(pretrained=True)  5vgg16 = models.vgg16(pretrained=True)  6densenet = models.densenet161(pretrained=True)  7inception = models.inception_v3(pretrained=True)  8googlenet = models.googlenet(pretrained=True)  9shufflenet = models.shufflenet_v2_x1_0(pretrained=True) 10mobilenet = models.mobilenet_v2(pretrained=True) 11resnext50_32x4d = models.resnext50_32x4d(pretrained=True) 12wide_resnet50_2 = models.wide_resnet50_2(pretrained=True) 13mnasnet = models.mnasnet1_0(pretrained=True)

这里，我们只需要执行resnet18 = models.resnet18(pretrained=True)就可以下载resnet18的模型。这些模型的输入格式要求如下：

大小都是224x224，

RGB三通道图像，

mean = [0.485, 0.456, 0.406]

std = [0.229, 0.224, 0.225]

下载与转为为ONNX的代码如下：

model = torchvision.models.resnet18(pretrained=True).eval() dummy_input = torch.randn((1, 3, 224, 224)) torch.onnx.export(model, dummy_input, "resnet18.onnx")

02

转为IR格式

Cmd至打开安装好的OpenVINO:

deployment_tools\model_optimizer

目录下，执行下面的命令行语句：

python mo_onnx.py --input_model D:\python\pytorch_tutorial\resnet18.onnx

可以看到resnet18模型已经成功转好！

03

OpenVINO SDK调用

对转换好的IR模型，就可以首先通过OpenVINO202R3的Python版本SDK完成加速推理预测，完整的代码实现如下：

from __future__ import print_function import cv2 import numpy as np import logging as log from openvino.inference_engine import IECore with open('imagenet_classes.txt') as f:     labels = [line.strip() for line in f.readlines()] def image_classification():     model_xml = "resnet18.xml"     model_bin = "resnet18.bin"     # Plugin initialization for specified device and load extensions library if specified     log.info("Creating Inference Engine")     ie = IECore()     # Read IR     log.info("Loading network files:\n\t{}\n\t{}".format(model_xml, model_bin))     net = ie.read_network(model=model_xml, weights=model_bin)     log.info("Preparing input blobs")     input_blob = next(iter(net.inputs))     out_blob = next(iter(net.outputs))     # Read and pre-process input images     n, c, h, w = net.inputs[input_blob].shape     images = np.ndarray(shape=(n, c, h, w))     src = cv2.imread("D:/images/messi.jpg")     image = cv2.resize(src, (w, h))     image = np.float32(image) / 255.0     image[:, :, ] -= (np.float32(0.485), np.float32(0.456), np.float32(0.406))     image[:, :, ] /= (np.float32(0.229), np.float32(0.224), np.float32(0.225))     image = image.transpose((2, 0, 1))     # Loading model to the plugin     log.info("Loading model to the plugin")     exec_net = ie.load_network(network=net, device_name="CPU")     # Start sync inference     log.info("Starting inference in synchronous mode")     res = exec_net.infer(inputs={input_blob: [image]})     # Processing output blob     log.info("Processing output blob")     res = res[out_blob]     label_index = np.argmax(res, 1)     label_txt = labels[label_index[0]]     cv2.putText(src, label_txt, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255), 2, 8)     cv2.imshow("ResNet18-from Pytorch image classification", src)     cv2.waitKey(0)     cv2.destroyAllWindows() if __name__ == '__main__':     image_classification()

运行结果如下：

善始者实繁

克终者盖寡

 推荐阅读 

OpenCV4系统化学习路线图-视频版本！

OpenCV单应性矩阵发现参数估算方法详解

单应性矩阵应用-基于特征的图像拼接

OpenCV图像拼接改进算法之完美拼接

OpenCV | 二值图像分析的技巧都在这里

OpenCV二值图像分析之形态学应用技巧

图像色彩空间与应用转换

五分钟学会C++高效图表绘制神器调用

没想到图像直方图有这么多应用场景

基于灰度共生矩阵(GLCM)的图像纹理分析与提取

OpenCV中一个最容易搞错的形态学操作

OpenCV实现皮肤表面粗糙度3D显示

解密 | OpenCV加载图像大小是有限制的 ？

OpenCV中ORB特征提取与匹配

OpenCV SIFT特征算法详解与使用

HOG特征详解与行人检测