4.6.2 构建onnx结构模型-特殊情况下的Gather

前言

构建onnx方式通常有两种：
1、通过代码转换成onnx结构，比如pytorch —> onnx
2、通过onnx 自定义结点，图，生成onnx结构

本文主要是简单学习和使用两种不同onnx结构，
下面以gather 结点进行分析

采用onnx方式构建

from onnx import TensorProto, helper
import onnx
import onnxruntime
import numpy as np

def construct_model_gather():
    print("construct model start... \n")
    input_shape1 = []
    indices = 2
    axis= 0
    input_shape0 = [3]
    output_shape = []

    initializer = [
        helper.make_tensor("indices", TensorProto.INT64, input_shape1, [indices])
    ]
    Gather = helper.make_node(
        "Gather", inputs=["data", "indices"], outputs=["output"], name="Gather_test"
    )
    Gather.attribute.extend([onnx.helper.make_attribute("axis", axis)])
    graph = helper.make_graph(
        nodes=[Gather],
        name="test_graph",
        inputs=[
            helper.make_tensor_value_info(
                "data", TensorProto.INT64, tuple(input_shape0)
            )
        ],  # use your input
        outputs=[
            helper.make_tensor_value_info(
                "output", TensorProto.INT64, tuple(output_shape)
            )
        ],
        initializer=initializer,
    )
    opset_imports = [onnx.helper.make_operatorsetid("", version=12)]
    model = onnx.helper.make_model(graph, opset_imports=opset_imports)
    model.ir_version = onnx.IR_VERSION
    onnx.checker.check_model(model)
    print("construct model done... \n")
    return model

def run(model_file):
    print(f'run start....\n')
    session = onnxruntime.InferenceSession(model_file, providers=['CPUExecutionProvider'])
    input_name1 = session.get_inputs()[0].name  
    input_data1= np.random.randn(3).astype(np.int64)
    print(f'input_data1:{input_data1}')

    pred_onx = session.run( None, {input_name1: input_data1})

    print(f'pred_onx:{pred_onx}\n')
    print(f'pred_onx[0].shape:{pred_onx[0].shape}\n')
    print(f'run done....\n')


if __name__ == '__main__':
    model = construct_model_gather()
    model_file = "test.onnx"
    onnx.save(model, model_file)
    run(model_file)