百度大多推的是对固定维度的模型进行简化
这里记录一下,如何对动态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!')