修改onnx模型中间节点命名(包含输入、输出重命名)

来源:Paddle2ONNX

Paddle2ONNX/tools/onnx/README.md at develop · PaddlePaddle/Paddle2ONNX · GitHub

依赖:import onnx

python rename_onnx_model.py --model model.onnx --origin_names x y z --new_names x1 y1 z1 --save_file new_model.onnx

其中 origin_names 和 new_names,前者表示原模型中各个命名(可指定多个),后者表示新命名,两个参数指定的命名个数需要相同 

 

import argparse
import sys


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--model',
        required=True,
        help='Path of directory saved the input model.')
    parser.add_argument(
        '--origin_names',
        required=True,
        nargs='+',
        help='The original name you want to modify.')
    parser.add_argument(
        '--new_names',
        required=True,
        nargs='+',
        help='The new name you want change to, the number of new_names should be same with the number of origin_names'
    )
    parser.add_argument(
        '--save_file', required=True, help='Path to save the new onnx model.')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_arguments()
    import onnx
    model = onnx.load(args.model)
    output_tensor_names = set()
    for ipt in model.graph.input:
        output_tensor_names.add(ipt.name)
    for node in model.graph.node:
        for out in node.output:
            output_tensor_names.add(out)

    for origin_name in args.origin_names:
        if origin_name not in output_tensor_names:
            print("[ERROR] Cannot find tensor name '{}' in onnx model graph.".
                  format(origin_name))
            sys.exit(-1)
    if len(set(args.origin_names)) < len(args.origin_names):
        print(
            "[ERROR] There's dumplicate name in --origin_names, which is not allowed."
        )
        sys.exit(-1)
    if len(args.new_names) != len(args.origin_names):
        print(
            "[ERROR] Number of --new_names must be same with the number of --origin_names."
        )
        sys.exit(-1)
    if len(set(args.new_names)) < len(args.new_names):
        print(
            "[ERROR] There's dumplicate name in --new_names, which is not allowed."
        )
        sys.exit(-1)
    for new_name in args.new_names:
        if new_name in output_tensor_names:
            print(
                "[ERROR] The defined new_name '{}' is already exist in the onnx model, which is not allowed."
            )
            sys.exit(-1)

    for i, ipt in enumerate(model.graph.input):
        if ipt.name in args.origin_names:
            idx = args.origin_names.index(ipt.name)
            model.graph.input[i].name = args.new_names[idx]

    for i, node in enumerate(model.graph.node):
        for j, ipt in enumerate(node.input):
            if ipt in args.origin_names:
                idx = args.origin_names.index(ipt)
                model.graph.node[i].input[j] = args.new_names[idx]
        for j, out in enumerate(node.output):
            if out in args.origin_names:
                idx = args.origin_names.index(out)
                model.graph.node[i].output[j] = args.new_names[idx]

    for i, out in enumerate(model.graph.output):
        if out.name in args.origin_names:
            idx = args.origin_names.index(out.name)
            model.graph.output[i].name = args.new_names[idx]

    onnx.checker.check_model(model)
    onnx.save(model, args.save_file)
    print("[Finished] The new model saved in {}.".format(args.save_file))
    print("[DEBUG INFO] The inputs of new model: {}".format(
        [x.name for x in model.graph.input]))
    print("[DEBUG INFO] The outputs of new model: {}".format(
        [x.name for x in model.graph.output]))
修改前

修改后
### 修改 ONNX 模型输入输出维度 对于需要调整 ONNX 模型输入输出维度的需求,可以采用多种方式来实现这一目标。具体的方法取决于模型的具体架构以及所使用的工具。 #### 方法一:通过修改 ONNX 文件中的定义 可以直接编辑 ONNX 文件中关于输入张量的信息来进行维度变更。例如,在处理行为识别模型时,如果原始模型接受五维输入,但目标平台 RKNN 仅支持四维,则可以通过重新定义 `input` 来适配新需求: ```python import onnx from onnx import helper, TensorProto # 加载现有的 ONNX 模型文件 model = onnx.load('original_model.onnx') # 创建新的输入描述符,这里假设将第一个维度设为动态(?),其他保持不变 new_input = helper.make_tensor_value_info( name='input', elem_type=TensorProto.FLOAT, shape=[None, 3, 256, 256] # 将原本的第一个维度改为 None 表示可变长度 ) # 替换旧的输入描述符 model.graph.input[0].CopyFrom(new_input) # 保存更改后的模型到磁盘 onnx.save(model, 'modified_model.onnx') ``` 这种方法适用于当只需要改变最外层的输入形状而不影响内部计算逻辑的情况下[^1]。 #### 方法二:利用 PyTorch 导出带有灵活维度的 ONNX 模型 另一种更为通用的方式是在最初从框架如 PyTorch 转换成 ONNX 的时候就设置好允许变化的维度参数。这使得即使后续操作会涉及到固定的中间节点尺寸,也能确保整个流程能够适应不同的输入规格: ```python import torch import torchvision.models as models import onnx dummy_input = torch.randn(1, 3, 256, 256) # 构建虚拟输入用于导出过程 model = models.resnet18(pretrained=True).eval() torch.onnx.export( model=model, args=(dummy_input,), f="dynamic_shape_model.onnx", input_names=['input'], output_names=['output'], dynamic_axes={ 'input': {0: 'batch_size'}, # 设置 batch 维度为动态 'output': {0: 'batch_size'} } ) ``` 此代码片段展示了如何创建一个具有动态批次大小的支持多批处理能力的 ONNX 模型[^2]。 #### 方法三:应用 Resize Operator 进行空间变换 有时可能不仅限于简单地调整时间序列或者批量数据的数量,还需要对图像或其他类型的特征图进行缩放。这时就可以考虑使用 ONNX 中内置的操作算子——Resize operator 来完成更复杂的重映射工作: ```python import numpy as np import onnxruntime as ort import onnx.helper as helper resize_node = helper.make_node( op_type="Resize", # 使用 Resize 算法 inputs=["X"], # 输入张量名称 outputs=["Y"] # 输出张量名称 ) scale_factor = [1., 1., 0.5, 0.5] # 定义缩小比例因子 roi = [] # ROI 参数为空列表表示全局范围内的 resize sizes = [] graph_def = helper.make_graph([resize_node], "test_resize", [], []) opset_imports = [helper.make_opsetid("", version=9)] model_def = helper.make_model(graph_def, producer_name="example", opset_imports=opset_imports) ort_session = ort.InferenceSession(onnx._serialize(model_def)) outputs = ort_session.run(None, {"X": np.random.rand(*input_dims)}) print(outputs['Y'].shape) # 打印经过 resize 后的结果形状 ``` 上述例子说明了怎样构建并执行一次基于指定尺度因子的空间变换操作[^3]。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值