Model Optimization——常用深度学习模型转换脚本

简介

​ 我们都知道目前流行的深度学习框架有很多种，例如Pytorch、TensorFlow、Caffe等，在实际训练和部署的过程中，常常会接触到不同框架模型之间的转换，大多数时候这些并不复杂，只需要写一些简短的脚本就能够完成。而像OpenVINO中IR这样的特殊模型，我们也可以通过其提供的自带脚本进行转换，下面简单介绍常用的模型之间转换的脚本和方法。

Pth模型 转 ONNX模型

import torch
import argparse

parser = argparse.ArgumentParser()


# the function about making the ".pth" model trained by pytorch convert to the "ONNX" model
def convert_model_to_ONNX(input_pth_model, output_ONNX_model, model_name):
    dummy_input = torch.randn(1, 1, 64, 64)

    model = torch.load(input_pth_model, map_location='cpu')

    input_names = ["input"]
    output_names = ["output"]

    model.eval()

    torch.onnx.export(model, dummy_input, output_ONNX_model + '/' + model_name, verbose=True, input_names=input_names,output_names=output_names)


if __name__ == '__main__':
    parser.add_argument("--input", type=str, help='input pth model path')
    parser.add_argument("--output", type=str, help='output onnx model path')
    parser.add_argument("--name", type=str, help='output onnx model filename')
    opt = parser.parse_args()
    # call the function named "convert_model_to_ONNX".
    convert_model_to_ONNX(opt.input, opt.output, opt.name)

Pth模型 转 Pt模型

import torch
from collections import OrderedDict
import pickle
import os
import sys

sys.path.insert(0, 'model dir')
device = 'cuda' if torch.cuda.is_available() else 'cpu'


if __name__ == '__main__':
    modelfile = 'yolov5s.pt'
    utl_model = torch.load('your output dir/'+modelfile)
    utl_param = utl_model['model'].model
    torch.save(utl_param.state_dict(), os.path.splitext(
        modelfile)[0] + '_param.pth')
    own_state = utl_param.state_dict()
    print(len(own_state))

    numpy_param = OrderedDict()
    for name in own_state:
        numpy_param[name] = own_state[name].data.cpu().numpy()
    print(len(numpy_param))
    with open(os.path.splitext(modelfile)[0] + '_numpy_param.pkl', 'wb') as fw:
        pickle.dump(numpy_param, fw)

Pt模型 转 ONNX模型

​ 使用YOLOv5源码中的export.py脚本，可以很容易地将yolo训练出的pt模型转化成通用的onnx模型。

​ 在yolov5的models文件夹下

python3 export.py --weights 模型路径 --img-size 输入层尺寸

​ 即可直接将pt模型转化成onnx模型。

ONNX模型 转 IR模型

​ OpenVINO推理加速使用的是一种中间表示(IR)模型，其自带的模型优化器中也提供了ONNX模型转换成IR模型的脚本，在deployment_tools/model_optimizer路径下

运行模型转换脚本

cd <INSTALL_DIR>/deployment_tools/model_optimizer/
python3 mo.py --input_model <INPUT_MODEL>.onnx  --output_dir <OUTPUT_MODEL_DIR> --model_name <OUTPUT_MODEL_NAME>

​ 出现SUCCESS则表示转换成功，在对应的路径下会生成.xml和.bin文件，在任何时候，你都需要将这两个模型文件放在统一路径下。

参考

Converting a Model to Intermediate Representation (IR)

后续

 喜欢的话可以关注一下我的公众号技术开发小圈，尤其是对深度学习以及计算机视觉有兴趣的朋友，我会把相关的源码以及更多资料发在上面，希望可以帮助到新入门的大家！