onnx模型结构优化相关记录

项目场景：

当我们使用torch模型转换为onnx模型时往往会遇到，在torch代码里就是一个算子，但是转换成onnx时就可能是好几个算子的组合形式。比如groupnorm算子就是一个很好的例子，转成onnx格式时如下所示：

问题描述

如下代码所示，我们使用torch定义了一个groupnorm算子，然后将该算子转换为onnx格式。

import torch
import torch.nn as nn

class PureGroupNormModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 创建一个只有GroupNorm层的模型
        self.gn = nn.GroupNorm(num_groups=4, num_channels=8)
    
    def forward(self, x):
        # 直接返回GroupNorm计算结果
        return self.gn(x)

# 创建模型实例
model = PureGroupNormModel()
model.eval()  # 设置为评估模式

# 创建符合GroupNorm要求的输入张量
# 格式: (batch_size, num_channels, height, width)
dummy_input = torch.randn(1, 8, 64, 64)  # 确保通道数匹配num_channels

# 导出ONNX文件
onnx_path = "pure_groupnorm.onnx"

torch.onnx.export(
    model,                   # 待导出的模型
    dummy_input,             # 模型输入
    onnx_path,              # 保存路径
    opset_version=13,        # ONNX算子集版本
    do_constant_folding=True,  # 优化常量
    input_names=['input'],   # 输入名称
    output_names=['output'], # 输出名称
)

print(f"Successfully exported ONNX model to {onnx_path}")

但是，我们在打开onnx文件的时候，却发现文件中并不是以groupnorm算子表示的，而是以几个小算子的组合形式出现。但是，这种算子组合可能会造成一个问题，就是推理框架在解析该onnx文件时，会以算子组合去跑。并不会调到groupnorm算子，影响推理性能。

因此，我们希望是转换成一个groupnorm算子的onnx文件。

原因分析：

首先，我们去查看一下ONNX中groupnorm的算子定义。发现groupnorm算子是在opset18版本下才开始支持。

下面是ONNX中groupnorm支持的版本：

我们在转换代码中，设置的opset=13，可能是这个原因造成的。因此，我们修改opset=18重新转换onnx模型文件。查看文件发现，依然是算子组合形式。

针对这个问题，询问一下deepseek解释一下原因：

其中，第一条和第二条是符合的。第三条确实没有设置，针对此问题重新设置导出：

依然失败。

解决方案：

在我实在想不到的情况下，发现了一个比较好用的工具就是：onnxruntime.transformers工具。

from onnxruntime.transformers import onnx_model
from onnxruntime.transformers.fusion import FusionGroupNorm

def fuse_groupnorm(onnx_path):
    # 加载模型
    model = onnx_model.OnnxModel(onnx.load(onnx_path))
    
    # 创建并应用融合器
    fusion_manager = fusion.FusionGroupNorm(model)
    fusion_manager.apply()
    
    # 保存优化后的模型
    model.save_model_to_file(onnx_path.replace('.onnx', '_fused.onnx'))

通过该方法成功的将groupnorm算子合并到一起，如下图所示：
而且该工具还默认将数据布局转换为NHWC格式。