如何將Pytorch自訓(xùn)練模型變成OpenVINO IR模型形式

簡 介

本文章將依次介紹如何將Pytorch自訓(xùn)練模型經(jīng)過一系列變換變成OpenVINO IR模型形式，而后使用OpenVINO Python API 對(duì)IR模型進(jìn)行推理，并將推理結(jié)果通過OpenCV API顯示在實(shí)時(shí)畫面上。

本文Python程序的開發(fā)環(huán)境是Ubuntu20.04 LTS + PyCharm，硬件平臺(tái)是AIxBoard愛克斯板開發(fā)者套件。

本文項(xiàng)目背景：針對(duì)2023第十一屆全國大學(xué)生光電設(shè)計(jì)競賽賽題2“迷宮尋寶”光電智能小車題目?；谠撡愴?xiàng)寶藏樣式，我通過深度學(xué)習(xí)訓(xùn)練出能分類四種不同顏色不同標(biāo)記形狀骨牌的模型，骨牌樣式詳見圖1.1。



| 圖1.1 四種骨牌類型

Pytorch pth模型轉(zhuǎn)換成OpenVINO IR模型

Pytorch是一個(gè)基于Torch的開源Python學(xué)習(xí)庫，是一個(gè)以Python優(yōu)先的深度學(xué)習(xí)框架。Pth模型文件是Pytorch進(jìn)行模型保存時(shí)的一種模型格式，OpenVINO暫不支持直接對(duì)Pth模型文件進(jìn)行推理，所以我們要將Pth格式的模型先轉(zhuǎn)換成ONNX格式文件，再通過OpenVINO自帶的Model Optimizer（模型優(yōu)化器）進(jìn)一步轉(zhuǎn)變成OpenVINO IR模型。處理過程如下所示：

通過Pytorch將Pth模型轉(zhuǎn)換成ONNX模型

轉(zhuǎn)換后的文件（Pth —> ONNX)：

import torch.onnx

# SZTU LIXROBO 23.5.14 #

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

# 1. 模型加載

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

# 2. 設(shè)置模型為評(píng)估模式而非訓(xùn)練模式

model.eval()

# 3. 生成隨機(jī)從標(biāo)準(zhǔn)正態(tài)分布抽取的張量

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

# 4. 導(dǎo)出ONNX模型（保存訓(xùn)練參數(shù)權(quán)重）

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

通過終端來將ONNX模型轉(zhuǎn)化成OpenVINO IR模型格式

在終端中輸入（Terminal）：

mo --input_model Domino_best.onnx --compress_to_fp16

# mo 啟動(dòng)OpenVINO 的Model Optimizer（模型優(yōu)化器） # input_model 輸入您轉(zhuǎn)換的ONNX模型內(nèi)容根的路徑 # compress _to_fp16 將模型輸出精度變?yōu)镕P16 等后一會(huì)，終端輸出：

?

代表ONNX模型轉(zhuǎn)換成OpenVINO IR模型成功。這里的信息告訴我們該Model是IR 11的形式，并分別保存在.xml和.bin文件下。

轉(zhuǎn)換后的文件（ONNX —> IR 11）：

?

mapping文件是一些轉(zhuǎn)換信息，暫時(shí)不會(huì)用到該文件。

至此，我們模型轉(zhuǎn)換的全部工作已經(jīng)完成，接下來就是運(yùn)用OpenVINO Runtime對(duì)IR 11模型進(jìn)行推理。

使用OpenVINO Runtime對(duì)IR 11模型進(jìn)行推理

在這一章節(jié)里我們將在Pycharm中使用OpenVINO Runtime對(duì)我們在1.2章中轉(zhuǎn)換得來的IR 11模型進(jìn)行推理，并將推理結(jié)果實(shí)時(shí)展現(xiàn)在攝像頭畫面中。

在開始之前，我們不妨了解推理程序的整個(gè)工作流程：

導(dǎo)入必要的功能庫（如openvino.runtime 以及 cv2和numpy）

探測硬件平臺(tái)所能使用的可搭載設(shè)備

創(chuàng)建核心對(duì)象以及加載模型和標(biāo)簽

輸入圖像進(jìn)行預(yù)處理，正則化，轉(zhuǎn)變成網(wǎng)絡(luò)輸入形狀

將處理后的圖像交由推理程序進(jìn)行推理，得到推理結(jié)果和處理時(shí)間并顯示出來

導(dǎo)入功能包

import openvino.runtime as ov

import numpy as np

import cv2

import time

這里一共導(dǎo)入4個(gè)功能包：

1. openvino.runtime 這是openvino runtime推理的主要功能包，也可用openvino.inference_engine進(jìn)行推理，過程大體是一致的。

2.numpy 這是常用的一個(gè)Python開源科學(xué)計(jì)算庫

3.cv2 也即OpenCV，用來處理有關(guān)圖像的一些信息

4.time 記錄系統(tǒng)運(yùn)行時(shí)間

設(shè)備檢測以及模型載入

我們可以使用Core對(duì)象中的available_devices函數(shù)來獲取當(dāng)前硬件平臺(tái)可供推理引擎使用的設(shè)備。

core = ov.Core()

print(core.available_devices)

如圖所示我們能得到在AlxBoard愛克斯開發(fā)板上可供我們使用的推理設(shè)備有CPU和GPU。

?

將模型進(jìn)行載入：

# SZTU LIXROBO 23.5.19 #

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

# 1. 創(chuàng)建核心對(duì)象

core = ov.Core()

# 2. 規(guī)定IR 11模型的模型文件和權(quán)重文件

model = "Domino_best.xml"

weights = "Domino_best.bin"

# 3. 將模型文件和權(quán)重文件進(jìn)行讀取

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

# 4. 把模型加載到設(shè)備上

（此處使用HETERO插件進(jìn)行異構(gòu)，加載到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手動(dòng)注入，也可使用導(dǎo)入標(biāo)簽文件等其他方式

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

圖像預(yù)處理

得到的圖像我們需要做一些預(yù)先處理才能輸入到推理引擎中進(jìn)行推理并得到結(jié)果。這一小節(jié)我們將展示如何把圖像進(jìn)行處理。

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

# 圖像預(yù)處理、歸一化 #

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

# 1. 類型轉(zhuǎn)換成np.float32

img = img.astype(np.float32)

# 2. 設(shè)置常用均值和標(biāo)準(zhǔn)差來正則化

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

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

# 圖像處理函數(shù) #

def img_pre(img):

# 1. 對(duì)OV輸入圖像顏色模型從BGR轉(zhuǎn)變成RGB

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

# 2. 對(duì)圖像進(jìn)行裁切

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

# 3. 使用我們定義的預(yù)處理函數(shù)對(duì)圖像進(jìn)行處理

nor_img = normalize(res_img)

# 4. 將處理好的圖像轉(zhuǎn)變?yōu)榫W(wǎng)絡(luò)輸入形狀

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

# 5. 返回處理結(jié)果

return nor_input_img

推理過程以及結(jié)果展示

在上一節(jié)中我們把輸入圖像所要進(jìn)行的預(yù)處理圖像進(jìn)行了一個(gè)定義，在這一小節(jié)則是OpenVINO Runtime推理程序的核心。

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

# 推理主程序 #

def image_infer(img):

# 1. 設(shè)置記錄起始時(shí)間

start_time = time.time()

# 2. 將圖像進(jìn)行處理

imgb = img_pre(img)

# 3. 輸入圖像進(jìn)行推理，得到推理結(jié)果

res_ir = com_model_ir([imgb])[output_layer_ir]

# 4. 對(duì)結(jié)果進(jìn)行歸一化處理，使用Sigmod歸一

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

# 5. 將結(jié)果進(jìn)行從小到大的排序，便于我們獲取置信度最高的類別

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

# 6. 用CV2的putText方法將置信度最高對(duì)應(yīng)的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. 記錄推理結(jié)束時(shí)間

end_time = time.time()

# 8. 計(jì)算出攝像頭運(yùn)行幀數(shù)

FPS = 1 / (end_time - start_time)

# 9. 將幀數(shù)繪制在圖像上

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

以上推理函數(shù)編寫已經(jīng)完成。以下是運(yùn)行主程序：

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

# 1. 獲取攝像頭

cap = cv2.VideoCapture(0)

# 2. 循環(huán)判斷

while 1:

# 1. 獲得實(shí)時(shí)畫面

success,frame = cap.read()

# 2. 把實(shí)時(shí)畫面交由推理函數(shù)進(jìn)行推理

frame = image_infer(frame)

# 3. 將畫面顯示在窗口

cv2.imshow("img",frame)

cv2.waitKey(1)

當(dāng)我們運(yùn)行該程序時(shí)，會(huì)得到如下畫面。



如圖所示，我們的Pytorch模型成功在OpenVINO的優(yōu)化以及推理下成功部署在AlxBoard愛克斯開發(fā)板，幀數(shù)在40-60之間，推理的結(jié)果非常好，很穩(wěn)定。

與Pytorch模型 CPU推理進(jìn)行比較

原先推理的過程我們是通過torch功能庫進(jìn)行推理，我們將兩者進(jìn)行比較。



| 左為OpenVINO優(yōu)化推理，右為torch推理

如圖所示OpenVINO優(yōu)化推理過后的結(jié)果從實(shí)際幀數(shù)上看大約有5-8倍的提升，推理精度也有少許加強(qiáng)。

結(jié) 論

自訓(xùn)練Pytorch模型在通過OpenVINO Model Optimizer 模型優(yōu)化后用OpenVINO Runtime進(jìn)行推理，推理過程簡單清晰。推理僅需幾個(gè)核心函數(shù)便可實(shí)現(xiàn)基于自訓(xùn)練Pytorch模型的轉(zhuǎn)化以及推理程序。OpenVINO簡單易上手，提供了強(qiáng)大的資料庫供學(xué)者查閱，其包含了從模型建立到模型推理的全過程。

審核編輯：劉清