Pytorch模型转成ONNX和MNN
Pytorch模型转成ONNX和MNN
准备工作
一、MNN安装
简介：MNN是一个轻量级的神经网络推理引擎，在端侧加载深度神经网络进行推理预测。由阿里巴巴团队研发并开源，目前在一些场景下性能和NCNN相同甚至超越；
-
下载MNN
git clone https://github.com/alibaba/MNN -
编译安装MNN
cd /path/to/MNN ./schema/generate.sh ./tools/script/get_model.sh # 可选，模型仅demo工程需要 mkdir build && cd build && cmake .. && make -j8
二、出现的问题
若出现没有安装protobuf的情况，需要安装protobuf
1、下载地址
https://github.com/protocolbuffers/protobuf/releases网站下，下载protobuf-all-3.17.3.tar.gz文件
2、解压
tar -xzvf protobuf-cpp-3.7.1.tar.gz
3、编译安装
cd protobuf-3.7.1
./autogen.sh
./configure
make
make check
sudo make install
实战
本次实战以mobilenet为例，需要先下载模型权重以及mobilenet.py文件
权重下载地址：'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
一、将pth文件转为onnx
1、需先安装torch和onnx库
pip install onnx
2、实现转换的代码：
import torch
import torch.nn as nn
import torch.onnx
import onnx
from mobilenet import mobilenet_v2 # 该文件为模型结构文件
pt_model_path = './mobilenet_v2-b0353104.pth'
onnx_model_path = './mobilenet_v2-b0353104.onnx'
model = mobilenet_v2(pretrained=False)
model.load_state_dict(torch.load(pt_model_path, map_location=torch.device('cpu')))
input_tensor = torch.randn(1, 3, 224, 224)
input_names = ['input']
output_names = ['output']
torch.onnx.export(model, input_tensor, onnx_model_path, verbose=True, input_names=input_names, output_names=output_names)
二、ONNX转为MNN
执行以下命令：
../../mnn/build/MNNConvert -f ONNX --modelFile mobilenet_v2-b0353104.onnx --MNNModel mobilenet_v2-b0353104.mnn --bizCode MNN
重点是找到mnn/build文件下的转换模型的可执行文件！
运行成功结果如下：
三、测试MNN文件
1、需先安装MNN库
pip install MNN
若出现以下问题：
Exception: initMNN: init numpy failed
解决方案：
如果MNN版本大于等于1.1.3，则Numpy版本需大于等于1.20.0；
如果MNN版本小于等于1.1.0，则Numpy版本需小于1.19.0；
2、代码实现调用MNN模型推理
方案一：调用python中MNN库
官方代码如下：
import MNN.expr as F
from torchvision import transforms
from PIL import Image
mnn_model_path = './mobilenet_v2-b0353104.mnn'
image_path = './test.jpg'
vars = F.load_as_dict(mnn_model_path)
inputVar = vars["input"]
# 查看输入信息
print('input shape: ', inputVar.shape)
# print(inputVar.data_format)
# 写入数据
input_image = Image.open(image_path)
preprocess = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
inputVar.write(input_tensor.tolist())
# 查看输出结果
outputVar = vars['output']
print('output shape: ', outputVar.shape)
# print(outputVar.read())
cls_id = F.argmax(outputVar, axis=1).read()
cls_probs = F.softmax(outputVar, axis=1).read()
print("cls id: ", cls_id)
print("cls prob: ", cls_probs[0, cls_id])
！！！在数据传输部分有问题，将数据类型从NC4HW4转成NCHW后，模型在推理部分出错，无法定位问题，可能是MNN 库有未解决的问题，所以放弃python实现；
方案二：C++实现
代码实现在mnn目录下的demo/exec，使用MobileNet的代码为pictureRecognition.cpp；
运行方式：
./pictureRecognition.out ../../moiblenet.mnn ../../testcat.jpg
注意文件路径的位置；
pictureRecognition.out在编译后的build文件夹下；
结果：
总结
1、深度学习轻量级框架对于Python的支持还不太完善，有些问题无法解决；
2、pth模型转ONNX模型需要知道模型的输入shape、输入节点、输出节点，和量化过程类似，并且模型结构需要先定义，相对于tensorflow来说不利于部署；
3、端侧部署需要掌握C++的基本语法，能够使用C++的数据类型以及标准库对数据流进行处理；