Plugins/Gnn/make_echo_edge_classifier.py

0001 #!/usr/bin/env python3
0002 """Write Tests/Data/echo_edge_classifier.onnx, used by OnnxEdgeClassifierTests.
0003
0004 The model ignores the graph and hands back the node features it was given,
0005 flattened row by row, as its (numEdges x 1) edge scores. Fed numNodes *
0006 numFeatures edges, its scores are then exactly the node feature matrix the
0007 classifier stage passed it, which shows which features reached the model, in
0008 which order and with which scale.
0009
0010 The edge index is still used, for the shape of the output, so that onnxruntime
0011 has no reason to drop the input.
0012 """
0013
0014 from pathlib import Path
0015
0016 import onnx
0017 from onnx import TensorProto, helper
0018
0019 nodes = [
0020     helper.make_node("Shape", ["edge_index"], ["edge_index_shape"]),
0021     helper.make_node("Slice", ["edge_index_shape", "one", "two"], ["num_edges"]),
0022     helper.make_node("Concat", ["num_edges", "one"], ["score_shape"], axis=0),
0023     helper.make_node("Reshape", ["x", "score_shape"], ["scores"]),
0024 ]
0025 initializers = [
0026     helper.make_tensor("one", TensorProto.INT64, [1], [1]),
0027     helper.make_tensor("two", TensorProto.INT64, [1], [2]),
0028 ]
0029 graph = helper.make_graph(
0030     nodes,
0031     "echo_edge_classifier",
0032     [
0033         helper.make_tensor_value_info("x", TensorProto.FLOAT, ["nodes", "features"]),
0034         helper.make_tensor_value_info("edge_index", TensorProto.INT64, [2, "edges"]),
0035     ],
0036     [helper.make_tensor_value_info("scores", TensorProto.FLOAT, ["edges", 1])],
0037     initializers,
0038 )
0039 # An old opset and IR version, so that any supported onnxruntime reads it
0040 model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
0041 model.ir_version = 7
0042 onnx.checker.check_model(model)
0043 onnx.save(
0044     model,
0045     Path(__file__).resolve().parents[3] / "Data" / "echo_edge_classifier.onnx",
0046 )