亲测有效 | OpenVINO支持ONNX格式文件直接读取了

博客专家认证

2020-10-12 11:39:21

01 函数支持
OpenVINO 2020R04版本支持ONNX格式文件的直接读取了，使用的函数还是跟之前读取IR文件函数保持一致，只是第二个参数缺省为空时候，就会尝试读取ONNX格式文件。相关的函数与参数解释如下：

CNNNetwork InferenceEngine::Core::ReadNetwork(

         const std::string &  modelPath,

         const std::string &  binPath = {}

) const

其中
modelPath 表示模型输入路径，（.xml 或者.onnx）
binPath 表示IR格式化数据路径(*.bin)，如果为空，尝试读取跟modelPath同名的bin文件，如果失败就会尝试直接加载IR文件。
02
ResNet模型转换
这里我使用pytorch中torchvision自带的ResNet18预训练模型，首先把它从pth转换为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中间文件格式，参见下面的链接，这里不再赘述
从Pytorch 的ONNX到OpenVINO中IR中间层

这样我们就得到了ONNX格式的模型了。该模型是在ImageNet数据集上训练生成的，支持1000个类别的图像分类，对图像预处理的参数支持与输入格式如下：
输入图像: HxW=224x224，
通道：RGB三通道图像，
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
对比测试
这里我们分别使用ResNet18网络的ONNX格式与IR格式在OpenVINO 2020R04 版本上进行测试，最终得到的时间比较如下：

可以看出直接读取ONNX方式执行的主要瓶颈居然在加载网络这步，对应的函数为：

ExecutableNetwork InferenceEngine::Core::LoadNetwork(

         const CNNNetwork &   network,

         const std::string &         deviceName,

         const std::map< std::string, std::string > &          config = {}

)

还好，当处理视频或者多次循环调用模型推理的时候，该函数属于初始化步骤，只会执行一次，不太会影响FPS。但是这个慢的有点离谱！希望可以在下个版本中可以改进一波。在推理执行上，两个格式基本速度相当，处理视频时候FPS也基本都很稳定。OpenVINO支持ONNX模型加载与推理稳了，这个将会是很多pytorch开发者的福音，CPU加速模型不是梦了。下面可以看一下我执行的推理时间与FPS：

这速度，还用说吗？靠谱！
测试代码
image_classification方法默认执行IR格式加载，修改参数为True执行ONNX格式推理，代码如下：

from __future__ import print_function

import cv2

import numpy as np

import time

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(use_onnx=False):

    model_xml = "resnet18.xml"

    model_bin = "resnet18.bin"

    onnx_model = "resnet18.onnx"



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

    inf_start = time.time()

    if use_onnx:

        # 直接使用ONNX格式加载

        net = ie.read_network(model=onnx_model)

    else:

        # IR 格式加载

        net = ie.read_network(model=model_xml, weights=model_bin)

    load_time = time.time() - inf_start

    print("read network time(ms) : %.3f"%(load_time*1000))



    log.info("Preparing input blobs")

    input_blob = next(iter(net.input_info))

    out_blob = next(iter(net.outputs))



    # Read and pre-process input images

    n, c, h, w = net.input_info[input_blob].input_data.shape



    src = cv2.imread("D:/images/messi.jpg")

    # image = cv2.dnn.blobFromImage(src, 0.00375, (w, h), (123.675, 116.28, 103.53), True)

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

    start_load = time.time()

    exec_net = ie.load_network(network=net, device_name="CPU")

    end_load = time.time() - start_load

    print("load  time(ms) : %.3f" % (end_load * 1000))



    # Start sync inference

    log.info("Starting inference in synchronous mode")

    inf_start1 = time.time()

    res = exec_net.infer(inputs={input_blob: [image]})

    inf_end1 = time.time() - inf_start1

    print("infer onnx as network time(ms) : %.3f" % (inf_end1 * 1000))



    # Processing output blob

    log.info("Processing output blob")

    res = res[out_blob]

    label_index = np.argmax(res, 1)

    label_txt = labels[label_index[0]]

    inf_end = time.time()

    det_time = inf_end - inf_start1

    inf_time_message = "Inference time: {:.3f} ms, FPS:{:.3f}".format(det_time * 1000, 1000 / (det_time * 1000 + 1))

    cv2.putText(src, label_txt, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255), 2, 8)

    cv2.putText(src, inf_time_message, (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2, 8)

    cv2.imshow("ResNet18-from Pytorch image classification", src)

    cv2.waitKey(0)

    cv2.destroyAllWindows()





if __name__ == '__main__':

    image_classification(True)