利用onnxsim对onnx模型进行简化

hjxu2016

已于 2022-10-14 11:45:42 修改

阅读量5.5k

点赞数 5

分类专栏：推理框架|OnnxRunTime 文章标签： python 深度学习开发语言

于 2022-10-11 16:18:22 首次发布

本文链接：https://blog.csdn.net/hjxu2016/article/details/127265957

版权

推理框架|OnnxRunTime 专栏收录该内容

3 篇文章

订阅专栏

百度大多推的是对固定维度的模型进行简化
这里记录一下，如何对动态batch的模型进行简化
模型简化的作用不再复述

一、简化固定维度的模型

import onnx
from onnxsim import simplify
 
ONNX_MODEL_PATH = './model.onnx'
ONNX_SIM_MODEL_PATH = './model_simple.onnx'
 
if __name__ == "__main__":
    onnx_model = onnx.load(ONNX_MODEL_PATH)
    onnx_sim_model, check = simplify(onnx_model)
    assert check, "Simplified ONNX model could not be validated"
    onnx.save(onnx_sim_model, ONNX_SIM_MODEL_PATH)
    print('ONNX file simplified!')

二、对动态batch的onnx模型进行简化

onnsim已经支持了动态batch模型的简化过程
一下两个参数可以设置动态batch模型

dynamic_input_shape
input_shapes

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""                  
*  * *** *  * *  *      
*  *  *   **  *  *             
****  *   **  *  *                 
*  *  *   **  *  *         
*  * **  *  * ****                

@File    : simpleOnnx.py
@Date    : 2022/10/11/011
@Require :
@Author  : https://blog.csdn.net/hjxu2016
@Funtion : 简化动态batch的onnx模型

"""

import onnx
from onnxsim import simplify

ONNX_MODEL_PATH = './model.onnx'
ONNX_SIM_MODEL_PATH = './model_simple.onnx'

if __name__ == "__main__":
    onnx_model = onnx.load(ONNX_MODEL_PATH)

    onnx_sim_model, check = simplify(onnx_model, dynamic_input_shape=True,
                                      input_shapes={"images":[1, 3, 224, 224]}) # 0.4版本，需要把1改成-1
    assert check, "Simplified ONNX model could not be validated"
    onnx.save(onnx_sim_model, ONNX_SIM_MODEL_PATH)
    print('ONNX file simplified!')