合并两个onnx文件到一个onnx

最新推荐文章于 2024-08-12 17:53:32 发布

AIcoderAR

最新推荐文章于 2024-08-12 17:53:32 发布

阅读量1.6k

点赞数

CC 4.0 BY-SA版权

文章标签：开发语言 python

本文链接：https://blog.youkuaiyun.com/weixin_42343812/article/details/130986855

该代码段展示了如何将两个ONNX模型（一个编码器和一个解码器）进行合并。首先，它遍历解码器的节点，修改其名称、输入和输出以避免与编码器冲突。接着，将解码器的初始化器添加到编码器中。然后，添加一个新的输出节点（ReduceSum操作）。最后，保存合并后的ONNX模型。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import onnx
import copy

def combine_onnx(list_modelname, output_combine_name):
    # 模型中节点和输入输出名字可能一样，合并前vocoder要加个前缀
    # onnx包含的数据类型： node\input\output\tensor(initializer)

    model_encoder = onnx.load(list_modelname[0])
    model_vocoder = onnx.load(list_modelname[1])
    #1、替换vocoder中的所有node的name、input、output并添加到encoder中
    for node in model_vocoder.graph.node:
        if 'c' in node.input:
            node.input.pop(0)
            node.input.append('dict')
        
        if node.name[0] == '/':
            node.name = '/vocoder_' + node.name[1:]

        for index in range(len(node.input)):
            if 'dict' in node.input: 
                input = node.input.pop(0)
                node.input.append('dict')

                continue
            input = node.input.pop(0)

            input = '/vocoder_' + input[1:] if input[0] == '/' else '/vocoder_' + input
            node.input.append(input)

        for index in range(len(node.output)):
            output = node.output.pop(0)
            output = '/vocoder_' + output[1:] if output[0] == '/' else '/vocoder_' + output
            node.output.append(output)

        model_encoder.graph.node.append(node)

    #2、添加vocoder中的initializer

    for weight in model_vocoder.graph.initializer:
        weight.name = '/vocoder_' + weight.name[1:] if weight.name[0] == '/' else '/vocoder_' + weight.name
        model_encoder.graph.initializer.append(weight)

    #3、添加 output
    model1_output_tensor_type = model_vocoder.graph.output[0].type.tensor_type
    output_elem_type = model1_output_tensor_type.elem_type
    output_shape = []
    output_name = '/vocoder_wav'

    for s in model1_output_tensor_type.shape.dim:
        if (s.dim_value > 0):
            output_shape.append(s.dim_value)
        else:        
            output_shape.append(s.dim_param)
    
    model_encoder_output = onnx.helper.make_tensor_value_info(
                                    output_name,
                                    output_elem_type,
                                    output_shape
                                )


    #4、添加duration sum 节点

    node = onnx.helper.make_node(
        'ReduceSum',
        name='output_duration',
        inputs=['duration'],
        outputs=['/output_duration_sum'],
    )
    node.attribute.append(onnx.helper.make_attribute("keepdims", 0))


    model_encoder.graph.node.append(node)

    model1_output_tensor_type = model_encoder.graph.output[1].type.tensor_type
    output_elem_type = model1_output_tensor_type.elem_type
    output_shape = [1]
    output_name = '/output_duration_sum'
    
    model_encoder_dur_output = onnx.helper.make_tensor_value_info(
                                    output_name,
                                    output_elem_type,
                                    output_shape
                                )
    model_encoder.graph.output.pop(0)
    model_encoder.graph.output.pop(0)
    model_encoder.graph.output.append(model_encoder_output)
    model_encoder.graph.output.append(model_encoder_dur_output)


    #4、保存新模型
    onnx.save(model_encoder, output_combine_name)

if __name__ == "__main__":
    list_modelname = ['fs2_conformer/encoder_ipu.onnx', 'fs2_conformer/vocoder.online.onnx']
    output_name = 'fs2_conformer/combine.onnx'
    combine_onnx(list_modelname, output_name)