ONNX模型优化：将卷积权重从Constant节点迁移到Initializer

背景介绍

在深度学习模型部署过程中，ONNX(Open Neural Network Exchange)格式已成为行业标准。ONNX模型使用图结构表示神经网络，其中包含两种主要的权重存储方式：Initializer(初始化器)和Constant(常量节点)。Initializer是ONNX推荐的标准方式，而Constant节点则会将权重作为计算图的一部分。

问题分析

在实际工作中，我们可能会遇到ONNX模型将卷积层权重存储在Constant节点而非Initializer中的情况。这种存储方式会导致模型图的initializer列表为空，可能影响后续的模型优化和部署流程。通过Netron可视化工具可以看到，这些权重确实存在于图中，但存储形式不符合最佳实践。

解决方案

我们可以使用ONNX Script库提供的IR(中间表示)接口来解决这个问题。具体步骤如下：

1. 加载模型：使用ONNX Script的IR模块加载原始模型
2. 识别卷积权重：遍历图中的Constant节点，识别出作为卷积权重的节点
3. 创建Initializer：将这些Constant节点的张量值提取出来，创建对应的Initializer
4. 替换引用：将图中所有对该Constant节点的引用替换为对新Initializer的引用
5. 清理节点：安全地移除原始的Constant节点
6. 保存模型：将优化后的模型保存到新文件

实现代码

from onnxscript import ir
from onnxscript.ir import convenience as ir_convenience

def convert_constants_to_initializers(model_path, output_path):
    # 加载原始模型
    model = ir.load(model_path)
    
    def is_conv_weight(node):
        # 检查节点输出是否被卷积层用作权重
        uses = {(node.op_type, idx) for node, idx in node.outputs[0].uses()}
        return ("Conv", 1) in uses
    
    # 遍历图中的所有节点
    for node in model.graph:
        if node.op_type == "Constant" and is_conv_weight(node):
            # 获取Constant节点的张量值
            value = node.outputs[0]
            tensor = node.attributes["value"].value
            
            # 创建对应的Initializer
            initializer_name = value.name
            assert initializer_name
            initializer = ir.Value(name=initializer_name, const_value=tensor)
            model.graph.initializers[initializer_name] = initializer
            
            # 替换所有引用
            ir_convenience.replace_all_uses_with(value, initializer)
            
            # 移除Constant节点
            model.graph.remove(node, safe=True)
    
    # 保存优化后的模型
    ir.save(model, output_path)

技术细节

1. 安全移除节点：使用safe=True参数确保只有在节点输出不再被使用时才移除
2. 引用替换：replace_all_uses_with函数会自动更新图中所有对该值的引用
3. Initializer存储：Initializer会被存储在模型的graph.initializers字典中，键为张量名称

应用场景

这种转换在以下场景特别有用：

• 准备模型进行量化时，需要所有权重都在Initializer中
• 优化模型图结构，减少不必要的计算节点
• 提高模型加载效率，Initializer比Constant节点更高效
• 统一模型格式，便于后续处理工具链兼容

注意事项

1. 执行转换前建议备份原始模型
2. 转换后应验证模型输出是否与原始模型一致
3. 对于大型模型，此操作可能会消耗较多内存
4. 确保ONNX Script版本与模型版本兼容

通过这种方法，我们可以有效地优化ONNX模型的结构，使其更符合最佳实践，为后续的部署和优化工作打下良好基础。