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