从Pytorch 的ONNX到OpenVINO中IR中间层

OpenCV学堂

发布于 2020-08-11 12:09:31

3.3K0

发布于 2020-08-11 12:09:31

微信公众号：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为例来演示。

下载模型与转ONNX格式

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

import torchvision.models as models
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
squeezenet = models.squeezenet1_0(pretrained=True)
vgg16 = models.vgg16(pretrained=True)
densenet = models.densenet161(pretrained=True)
inception = models.inception_v3(pretrained=True)
googlenet = models.googlenet(pretrained=True)
shufflenet = models.shufflenet_v2_x1_0(pretrained=True)
mobilenet = models.mobilenet_v2(pretrained=True)
resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
wide_resnet50_2 = models.wide_resnet50_2(pretrained=True)
mnasnet = 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")

转为IR格式

Cmd至打开安装好的OpenVINO:

deployment_tools\model_optimizer

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

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

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

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()

运行结果如下：