如何將Pytorch生成的模型進行CPU部署

我們使用Pytorch深度學習框架訓練好模型通常保存爲.pth文件,但這種網絡文件結構通常只在實驗驗證或者網絡學習時使用。如果你做好了一個網絡,想要將其部署在終端,或者想要加速其推理速度,可能需要其他專門的框架來幫助其部署。

本文介紹英特爾公司開發的OpenVINO工具箱,用其來進行Pytorch生成模型的部署和加速。

OpenVINO介紹

OpenVINO是英特爾公司開發的基於卷積神經網絡(CNN)的深度學習推理加速工具箱。它能使英特爾公司硬件,最大化提高深度學習計算性能。因此,當將深度學習模型部署到CPU端時,可以利用OpenVINO工具箱進行部署,提高深度學習模型推理速度。

OpenVINO的安裝可以見官網,從裏面就可以下載安裝。下面這是指向鏈接。需要注意的是,如果你在Windows平臺安裝的話,需要同時安裝Visual Studio2015或2017的c++環境,此外還需要安裝Python和Cmake工具。關於Python和Cmake都可以在對應官網上下載安裝。安裝Python和Cmake的時候,需要注意應當將其添加到環境變量中,否則可能在安裝OpenVINO是會出一些錯誤。

OpenVINO官網
OpenVINO安裝指導教程

ONNX

OpenVINO無法直接識別Pytorch生成的pt文件,需要先將pt文件轉成ONNX的格式。轉換示例代碼如下:

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# torch.load('filename.pth').to(device)

model = torch.load('ckpt_15.pth', map_location=device)
model.eval()
batch_size = 1  #批處理大小
input_shape = (3,320,320)   #輸入數據

input_data_shape = torch.randn(batch_size, *input_shape, device=device)

torch.onnx.export(model, input_data_shape, "filename.onnx", verbose=True)

需要注意的是,最好添加model.eval()這行語句,如果你的網絡在訓練時有Dropout,使用這個語句就能將一些值固定下來。

有一點需要特別強調的是,訓練時一定要保存模型本身文件而不是state_dict文件。加載state_dict文件的時候會報錯:

‘collections.OrderedDict’ object has no attribute ‘eval’

不是真正的模型文件是沒法進行調用.eval()方法的。

所以訓練階段對模型保存時要寫成這樣的語句:

torch.save(model, f"../output/ckpt_%d.pth" % epoch)

而不要寫成

torch.save(model.state_dict(), f"../output/ckpt_%d.pth" % epoch)

轉換成bin和xml

OpenVINO可以直接處理的網絡格式時它自有的bin和xml格式。對於已經獲得的ONNX文件,可以使用OpenVINO自帶的工具完成模型轉換。

首先:

Go to the<INSTALL_DIR>/deployment_tools/model_optimizerdirectory.

之後:

python3 mo.py --input_model <INPUT_MODEL>.onnx

即可完成模型轉換,生成bin和xml文件。

這裏,列出OpenVINO支持的網絡層轉換類型,如果你設計的網絡層裏面有OpenVINO不支持的網絡層,很有可能會轉換失敗。

OpenVINO支持的網絡層

參考資料

[1] collections.OrderedDict’ object has no attribute ‘eval’

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章