ONNX Simplifier: 简化您的ONNX模型

ONNX Simplifier: 让复杂的ONNX模型变得简单

ONNX (Open Neural Network Exchange) 是一种用于表示深度学习模型的开放格式,它允许AI开发者更轻松地在不同的框架和工具之间转换模型。然而,有时从框架导出的ONNX模型可能会变得非常复杂,包含许多冗余的操作。这就是ONNX Simplifier发挥作用的地方。

ONNX Simplifier的诞生背景

ONNX Simplifier的创建者在尝试将一个简单的reshape操作导出为ONNX模型时,发现生成的模型出乎意料地复杂。例如,下面这段代码:

import torch

class JustReshape(torch.nn.Module):
    def __init__(self):
        super(JustReshape, self).__init__()

    def forward(self, x):
        return x.view((x.shape[0], x.shape[1], x.shape[3], x.shape[2]))

net = JustReshape()
model_name = 'just_reshape.onnx'
dummy_input = torch.randn(2, 3, 4, 5)
torch.onnx.export(net, dummy_input, model_name, input_names=['input'], output_names=['output'])

理想情况下,这应该生成一个简单的模型,如下图所示:

简单的reshape模型

但实际上,生成的模型却复杂得多:

复杂的reshape模型

这种复杂性不仅降低了模型的可读性,还可能影响模型的执行效率。为了解决这个问题,ONNX Simplifier应运而生。

ONNX Simplifier的工作原理

ONNX Simplifier的核心思想是通过推理整个计算图,然后用常量输出替换冗余的操作符(也称为常量折叠)来简化ONNX模型。这个过程可以显著减少模型中的节点数量,使模型结构更加清晰,同时保持原始功能不变。

如何使用ONNX Simplifier

ONNX Simplifier提供了两种使用方式:Web版本和Python版本。

Web版本

对于那些不想安装任何软件的用户,ONNX Simplifier提供了一个在线版本,可以在convertmodel.com上使用。这个版本完全在浏览器中本地运行,确保了模型的安全性。

Python版本

对于希望在本地或脚本中使用ONNX Simplifier的用户,可以通过pip安装Python包:

pip3 install -U pip && pip3 install onnxsim

安装完成后,可以使用命令行工具简化ONNX模型:

onnxsim input_onnx_model output_onnx_model

对于更高级的功能,可以使用onnxsim -h查看帮助信息。

在Python脚本中使用ONNX Simplifier

ONNX Simplifier也可以很容易地集成到Python脚本中:

import onnx
from onnxsim import simplify

# 加载预定义的ONNX模型
model = onnx.load(filename)

# 转换模型
model_simp, check = simplify(model)

assert check, "Simplified ONNX model could not be validated"

# 使用model_simp作为标准ONNX模型对象