我們使用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 torchdevice = 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’

總結

以上是生活随笔為你收集整理的如何将Pytorch生成的模型进行CPU部署的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。