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

簡 介

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

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

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



| 圖1.1 四種骨牌類型

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

Pytorch是一個基于Torch的開源Python學(xué)習(xí)庫，是一個以Python優(yōu)先的深度學(xué)習(xí)框架。Pth模型文件是Pytorch進(jìn)行模型保存時的一種模型格式，OpenVINO暫不支持直接對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è)置模型為評估模式而非訓(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 啟動OpenVINO 的Model Optimizer（模型優(yōu)化器） # input_model 輸入您轉(zhuǎn)換的ONNX模型內(nèi)容根的路徑 # compress _to_fp16 將模型輸出精度變?yōu)镕P16 等后一會，終端輸出：

?

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

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

?

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

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

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

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

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

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

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

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

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

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

導(dǎo)入功能包

import openvino.runtime as ov

import numpy as np

import cv2

import time

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

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

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

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

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

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

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

core = ov.Core()

print(core.available_devices)

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

?

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

# SZTU LIXROBO 23.5.19 #

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

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

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ǎ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. 對OV輸入圖像顏色模型從BGR轉(zhuǎn)變成RGB

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

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

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

# 3. 使用我們定義的預(yù)處理函數(shù)對圖像進(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)行了一個定義，在這一小節(jié)則是OpenVINO Runtime推理程序的核心。

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

# 推理主程序 #

def image_infer(img):

# 1. 設(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. 對結(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方法將置信度最高對應(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é)束時間

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í)時畫面

success,frame = cap.read()

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

frame = image_infer(frame)

# 3. 將畫面顯示在窗口

cv2.imshow("img",frame)

cv2.waitKey(1)

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



如圖所示，我們的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)行推理，推理過程簡單清晰。推理僅需幾個核心函數(shù)便可實(shí)現(xiàn)基于自訓(xùn)練Pytorch模型的轉(zhuǎn)化以及推理程序。OpenVINO簡單易上手，提供了強(qiáng)大的資料庫供學(xué)者查閱，其包含了從模型建立到模型推理的全過程。

審核編輯：劉清