我們使用Pytorch深度學習框架訓練好模型通常保存爲.pth文件,但這種網絡文件結構通常只在實驗驗證或者網絡學習時使用。如果你做好了一個網絡,想要將其部署在終端,或者想要加速其推理速度,可能需要其他專門的框架來幫助其部署。
本文介紹英特爾公司開發的OpenVINO工具箱,用其來進行Pytorch生成模型的部署和加速。
OpenVINO介紹
OpenVINO是英特爾公司開發的基於卷積神經網絡(CNN)的深度學習推理加速工具箱。它能使英特爾公司硬件,最大化提高深度學習計算性能。因此,當將深度學習模型部署到CPU端時,可以利用OpenVINO工具箱進行部署,提高深度學習模型推理速度。
OpenVINO的安裝可以見官網,從裏面就可以下載安裝。下面這是指向鏈接。需要注意的是,如果你在Windows平臺安裝的話,需要同時安裝Visual Studio2015或2017的c++環境,此外還需要安裝Python和Cmake工具。關於Python和Cmake都可以在對應官網上下載安裝。安裝Python和Cmake的時候,需要注意應當將其添加到環境變量中,否則可能在安裝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不支持的網絡層,很有可能會轉換失敗。