分享 | 将Pytorch模型部署到Movidius神经计算棒


内容提要

这篇文章将是从笔者一个小白的视角出发,演示如何将自己训练的网络模型使用OpenVINO的优化器进行优化,并将其部署到神经计算棒进行推理加速的过程。

正文部分

1.将Pytorch模型转化为ONNX格式

我们可以从脚本中保存网络的结构开始看,Pytorch保存模型分为两种:

  ① 只保留模型参数
只保留参数[/align]torch.save(model.state_dict(), '\parameter.pkl')

②保存完整的模型

保存完整的模型torch.save(model, '\model.pkl')

我们需要完整的网络信息用于转化,所以需要在此选择第二种“保存完整的模型”方式。

之后,我们需要读取这个 pkl 文件并将其转化为 onnx 格式,我们需要定义批处理大小、输入数据尺寸、使用设备等信息,再定义好文件输入与输出的路径即可。

pytorch 计算acc pytorch 计算棒_xml

2.使用模型优化器进行优化

OpenVINO 的推理引擎支持的文件为中间表示文件 IR ,需要调用 OpenVINO 工具包中的mo.py 进行优化,如果使用 Tensorflow 等直接支持的框架则可以直接从框架模型转化为 IR ,而如果是像本文中提到的 Pytorch 的模型,其需先转化到 ONNX 文件再进行间接的IR转化,最终得到的IR文件将由两个部分组成:

① .xml 描述网络的拓扑结构
② .bin 网络的权重与偏置的二进制数据

先在开始前设置好 OpenVINO 的环境变量。

pytorch 计算acc pytorch 计算棒_pytorch 计算acc_02

接着,我们需要切换到模型优化器的目录下,这个目录下保存了用于不同框架的模型优化器,例如 mo_caffe.py, mo_tf.py 等文件。

pytorch 计算acc pytorch 计算棒_神经网络_03

当然,如果安装时没有安装在默认目录的话,需要重新调整前半段的路径位置。

cd C:\Program Files
(x86)\IntelSWTools\openvino\deployment_tools\model_optimizer

模型优化后的文件会自动保存在【…\openvino\deployment_tools\model_optimizer\example.xml】目录下,我们需要记下这个位置之后在部署时使用。

3.将优化后的模型部署到神经网络计算棒上

笔者在比赛中的项目是射频信号图片的分类工作,需要用到 classification_sample.py ,在 model 参数中输入上一步生成的 xml 文件的路径,在 input 参数中输入需要进行分类的图片路径,我们也可以通过 device 参数选择运行的设备,如果需要使用神经计算棒进行加速的话,我们在这一栏选择 MYRIAD 即可。

由于 OpenVINO 在顶层规划上就很好地采用了一套 API 用于不同设备的设计,我们在实际使用的时候输入 xml 的拓扑结构描述文件与需要进行加速的设备就能方便地使用了,这也是其一大亮点。

"C:\Program Files (x86)\IntelSWTools\openvino_2020.3.194\deployment_tools\inference_engine\samples\python\classification_sample\classification_sample.py" --model 【xml_path】 --device HETERO:MYRIAD,CPU,GPU --input 【img_path】

pytorch 计算acc pytorch 计算棒_神经网络_04

从输出我们可以看到神经计算棒顺利的完成了这个四分类问题的推理,尽管经过了模型的转化和优化,但是推理的正确率仍然是很高的,神经计算棒的推理速度也非常快。当然,这只是一个 getting start 的一个简单示例,在实际生产环境中还是要使用专用的 benchmark_app 工具。