如何将Pytorch生成的模型进行CPU部署

OpenVINO 专栏收录该内容
2 篇文章 1 订阅

我们使用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’

  • 7
    点赞
  • 1
    评论
  • 25
    收藏
  • 一键三连
    一键三连
  • 扫一扫,分享海报

©️2021 CSDN 皮肤主题: 技术黑板 设计师:CSDN官方博客 返回首页
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值