在使用onnx模型在其他框架进行部署时,一般都需要对onnx模型进行简化和优化。

1 onnx模型简化

对模型进行简化一般使用onnx-simplifier

可以通过以下命令安装

pip install onnx-simplifier

然后就可以在项目导出onnx模型之后,使用以下代码对onnx模型进行简化

import onnx
from onnxsim import simplify

if __name__ == '__main__':
    onnx_path = './mobile.onnx'
    onnx_simp_path = './mobile_simp.onnx'
    onnx_model = onnx.load(onnx_path)  # 加载onnx模型
    onnx_model_simp, check = simplify(onnx_model)  # 简化onnx模型
    assert check, "Simplified ONNX model could not be validated"
    onnx.save(onnx_model_simp, onnx_simp_path)

也可以在python环境中使用以下命令行完成简化

python3 -m onnxsim ./mobile.onnx ./mobile_simp.onnx

2 onnx模型优化

对模型进行简化一般使用onnxoptimizer

可以通过以下命令安装

pip install onnxoptimizer

使用以下代码对onnx模型进行优化

import onnx
import onnxoptimizer

if __name__ == '__main__':
    onnx_path = './mobile.onnx'
    onnx_optim_path = './mobile_optim.onnx'
    onnx_model = onnx.load(onnx_path)  # 加载onnx模型
    onnx_model_optim = onnxoptimizer.optimize(onnx_model)  # 优化onnx模型
    onnx.save(onnx_model_optim, onnx_optim_path)