onnx文件修改输入并自动推理

这次在工作中遇到了这么一个问题，就是修改onnx的时候，只能对指定的层进行大小的修改，比如输入输出，但是保存的onnx文件的中间层无法进行自动更新，于是在查找资料之后，解决了这个问题，做出如下总结

import onnx
from onnx.shape_inference import infer_shapes
import tflite2onnx


def onnx_as(input_size, onnx_path, onnx_save):
    '''
    onnx文件参数的修改
    :param model_path: 需要修改的onnx文件地址
    :return:
    '''
    model = onnx.load_model(onnx_path)
    # 获取原模型图表信息
    graph = model.graph
    # 对输入进行修改
    d = model.graph.input[0].type.tensor_type.shape.dim
    d[2].dim_value = input_size[0]
    d[3].dim_value = input_size[1]
    
    output2Height = input_size[0]
    output2Width = input_size[1]
    
    output1Height = output2Height // 2
    output1Width = output2Width // 2
    
    output0Height = output1Height // 2
    output0Width = output1Width // 2
    
    def INSIZE2OUTSIZE(x):
        return ((x + 1) // 2 + 1) // 2 - 7    
    
    # 对输出进行修改
    output_0 = model.graph.output[0].type.tensor_type.shape.dim
    output_1 = model.graph.output[1].type.tensor_type.shape.dim
    output_2 = model.graph.output[2].type.tensor_type.shape.dim
    output_0[2].dim_value = INSIZE2OUTSIZE(output0Height)
    output_0[3].dim_value = INSIZE2OUTSIZE(output0Width)
    output_1[2].dim_value = INSIZE2OUTSIZE(output1Height)
    output_1[3].dim_value = INSIZE2OUTSIZE(output1Width)
    output_2[2].dim_value = INSIZE2OUTSIZE(output2Height)
    output_2[3].dim_value = INSIZE2OUTSIZE(output2Width)
    
    # onnx.helper.make_graph函数：构建计算图原型，需要传入节点、图名称、输入张量信息、输出张量
    graph = onnx.helper.make_graph(graph.node, graph.name, graph.input, graph.output, graph.initializer)
    # 把计算图封装进模型modelproto
    info_model = onnx.helper.make_model(graph)
    # 进行形状推断
    onnx_model = onnx.shape_inference.infer_shapes(info_model)
    # 进行模型检测
    onnx.checker.check_model(model)
    # 进行模型保存
    onnx.save_model(onnx_model, onnx_save)
    

if __name__ == '__main__':
    onnx_path = r"test/model.onnx"
    onnx_save = r"test/model_as.onnx"

    input_size = (504, 896)
    onnx_file_as = onnx_as(input_size, onnx_path, onnx_save)