如何将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’

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