如何将Pytorch自训练模型变成OpenVINO IR模型形式

简  介  

本文章将依次介绍如何将Pytorch自训练模型经过一系列变换变成OpenVINO IR模型形式，而后使用OpenVINO Python API 对IR模型进行推理，并将推理结果通过OpenCV API显示在实时画面上。  

本文Python程序的开发环境是Ubuntu20.04 LTS + PyCharm，硬件平台是AIxBoard爱克斯板开发者套件。  

本文项目背景：针对2023第十一届全国大学生光电设计竞赛赛题2“迷宫寻宝”光电智能小车题目。基于该赛项宝藏样式，我通过深度学习训练出能分类四种不同颜色不同标记形状骨牌的模型，骨牌样式详见图1.1。



| 图1.1  四种骨牌类型  

Pytorch pth模型转换成OpenVINO IR模型  

Pytorch是一个基于Torch的开源Python学习库，是一个以Python优先的深度学习框架。Pth模型文件是Pytorch进行模型保存时的一种模型格式，OpenVINO暂不支持直接对Pth模型文件进行推理，所以我们要将Pth格式的模型先转换成ONNX格式文件，再通过OpenVINO自带的Model Optimizer（模型优化器）进一步转变成OpenVINO IR模型。处理过程如下所示：  

通过Pytorch将Pth模型转换成ONNX模型

转换后的文件（Pth —> ONNX)：

import torch.onnx

#  SZTU LIXROBO 23.5.14 #

#******************************************#

#  1. 模型加载

model = torch.load('Domino_best.pth', map_location=torch.device('cpu'))

#  2. 设置模型为评估模式而非训练模式

model.eval()

#  3. 生成随机从标准正态分布抽取的张量

dummy_input = torch.randn(1,3,224,224,device='cpu')

#  4. 导出ONNX模型（保存训练参数权重）

torch.onnx.export(model,dummy_input,"Domino_best.onnx",export_params=True)

通过终端来将ONNX模型转化成OpenVINO IR模型格式

在终端中输入（Terminal）：

mo --input_model Domino_best.onnx --compress_to_fp16

  #  mo  启动OpenVINO 的Model Optimizer（模型优化器） #  input_model 输入您转换的ONNX模型内容根的路径 #  compress _to_fp16 将模型输出精度变为FP16   等后一会，终端输出：

 

代表ONNX模型转换成OpenVINO IR模型成功。这里的信息告诉我们该Model是IR 11的形式，并分别保存在.xml和.bin文件下。  

转换后的文件（ONNX —> IR 11）：

 

mapping文件是一些转换信息，暂时不会用到该文件。  

至此，我们模型转换的全部工作已经完成，接下来就是运用OpenVINO Runtime对IR 11模型进行推理。

使用OpenVINO Runtime对IR 11模型进行推理  

在这一章节里我们将在Pycharm中使用OpenVINO Runtime对我们在1.2章中转换得来的IR 11模型进行推理，并将推理结果实时展现在摄像头画面中。  

在开始之前，我们不妨了解推理程序的整个工作流程：

导入必要的功能库（如openvino.runtime 以及 cv2和numpy）

探测硬件平台所能使用的可搭载设备

创建核心对象以及加载模型和标签

输入图像进行预处理，正则化，转变成网络输入形状

将处理后的图像交由推理程序进行推理，得到推理结果和处理时间并显示出来

导入功能包

import openvino.runtime as ov

import numpy as np

import cv2

import time

这里一共导入4个功能包：

1. openvino.runtime  这是openvino runtime推理的主要功能包，也可用openvino.inference_engine进行推理，过程大体是一致的。

2. numpy  这是常用的一个Python开源科学计算库

3. cv2  也即OpenCV，用来处理有关图像的一些信息

4. time  记录系统运行时间  

设备检测以及模型载入

我们可以使用Core对象中的available_devices函数来获取当前硬件平台可供推理引擎使用的设备。

core = ov.Core()

print(core.available_devices)

如图所示我们能得到在AlxBoard爱克斯开发板上可供我们使用的推理设备有CPU和GPU。

 

将模型进行载入：

#  SZTU LIXROBO 23.5.19 #

#************************************#

#  1.  创建核心对象

core = ov.Core()

#  2.  规定IR 11模型的模型文件和权重文件

model = "Domino_best.xml"

weights = "Domino_best.bin"

#  3.  将模型文件和权重文件进行读取

model_ir = core.read_model(model= model,weights=weights)

#  4.  把模型加载到设备上

（此处使用HETERO插件进行异构，加载到GPU和CPU上）

com_model_ir= core.compile_model(model=model_ir,device_name="HETERO:GPU,CPU")

#  5.  获取模型输出层

output_layer_ir = com_model_ir.outputs[0]

#  6.  由于是简单模型，故label手动注入，也可使用导入标签文件等其他方式

label = ['BlueFake','BlueTrue','RedFake','RedTrue']

图像预处理

得到的图像我们需要做一些预先处理才能输入到推理引擎中进行推理并得到结果。这一小节我们将展示如何把图像进行处理。

#************************************#

#           图像预处理、归一化          #

def normalize(img: np.ndarray) ->np.ndarray:

    #  1.  类型转换成np.float32

    img = img.astype(np.float32)

    #  2.  设置常用均值和标准差来正则化

    mean  =(0.485,0.456,0.406)

    std   =(0.299,0.224,0.255)

    img /=255.0

    img -=mean

    img /=std

    #  3.  返回处理后的img

    return img

#************************************#

#              图像处理函数            #

def img_pre(img):

    #  1.  对OV输入图像颜色模型从BGR转变成RGB

    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

    #  2.  对图像进行裁切

    res_img = cv2.resize(img, (224, 224))

    #  3.  使用我们定义的预处理函数对图像进行处理

    nor_img = normalize(res_img)

    #  4.  将处理好的图像转变为网络输入形状

    nor_input_img = np.expand_dims(np.transpose(nor_img, (2, 0, 1)), 0)

    #  5.  返回处理结果

    return nor_input_img

推理过程以及结果展示

在上一节中我们把输入图像所要进行的预处理图像进行了一个定义，在这一小节则是OpenVINO Runtime推理程序的核心。

#************************************#

#               推理主程序             #

def image_infer(img):

    #  1.  设置记录起始时间

    start_time = time.time()

    #  2.  将图像进行处理

    imgb = img_pre(img)

    #  3.  输入图像进行推理，得到推理结果

    res_ir = com_model_ir([imgb])[output_layer_ir]

    #  4.  对结果进行归一化处理，使用Sigmod归一

    Confidence_Level = 1/(1+np.exp(-res_ir[0]))

    #  5.  将结果进行从小到大的排序，便于我们获取置信度最高的类别

    result_mask_ir = np.squeeze(np.argsort(res_ir, axis=1)).astype(np.uint8)

    #  6.  用CV2的putText方法将置信度最高对应的label以及其置信度绘制在图像上

    img = cv2.putText(img,str(label[result_mask_ir[3]])+'  '+ str(Confidence_Level[result_mask_ir[3]]),(50,80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2,cv2.LINE_AA)

    #  7.  记录推理结束时间

    end_time = time.time()

    #  8.  计算出摄像头运行帧数

    FPS = 1 / (end_time - start_time)

#  9.  将帧数绘制在图像上

    img = cv2.putText(img, 'FPS  ' + str(int(FPS)), (50, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2,cv2.LINE_AA)

    #  10. 返回图像

    return img

  以上推理函数编写已经完成。以下是运行主程序：

#********************主程序***********************#

#  1.  获取摄像头

cap = cv2.VideoCapture(0)

#  2.  循环判断

while 1:

    #  1.  获得实时画面

    success,frame = cap.read()

    #  2.  把实时画面交由推理函数进行推理

    frame = image_infer(frame)

    #  3.  将画面显示在窗口

    cv2.imshow("img",frame)

    cv2.waitKey(1)

  当我们运行该程序时，会得到如下画面。



如图所示，我们的Pytorch模型成功在OpenVINO的优化以及推理下成功部署在AlxBoard爱克斯开发板，帧数在40-60之间，推理的结果非常好，很稳定。  

与Pytorch模型 CPU推理进行比较  

原先推理的过程我们是通过torch功能库进行推理，我们将两者进行比较。



| 左为OpenVINO优化推理，右为torch推理  

如图所示OpenVINO优化推理过后的结果从实际帧数上看大约有5-8倍的提升，推理精度也有少许加强。  

结  论  

自训练Pytorch模型在通过OpenVINO Model Optimizer 模型优化后用OpenVINO Runtime进行推理，推理过程简单清晰。推理仅需几个核心函数便可实现基于自训练Pytorch模型的转化以及推理程序。OpenVINO简单易上手，提供了强大的资料库供学者查阅，其包含了从模型建立到模型推理的全过程。  

  



审核编辑：刘清