From 9e0f8f21bf3d5be1b3e6be22542dc949a3352aa9 Mon Sep 17 00:00:00 2001
From: YdrMaster <ydrml@hotmail.com>
Date: Tue, 21 Feb 2023 14:01:57 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E7=94=9F=E6=88=90=E6=A8=A1=E5=9E=8B?=
 =?UTF-8?q?=E5=AF=B9=E8=B1=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: YdrMaster <ydrml@hotmail.com>
---
 pyinfinitensor/src/pyinfinitensor/onnx.py | 35 ++++++++++++++++-------
 pyinfinitensor/tests/test_onnx.py         |  2 +-
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/pyinfinitensor/src/pyinfinitensor/onnx.py b/pyinfinitensor/src/pyinfinitensor/onnx.py
index 03339455..802fdb15 100644
--- a/pyinfinitensor/src/pyinfinitensor/onnx.py
+++ b/pyinfinitensor/src/pyinfinitensor/onnx.py
@@ -7,7 +7,13 @@ from onnx import (
     TensorShapeProto,
     ValueInfoProto,
 )
-from onnx.helper import make_node, make_tensor_value_info, make_tensor
+from onnx.helper import (
+    make_node,
+    make_tensor_value_info,
+    make_tensor,
+    make_graph,
+    make_model,
+)
 from onnx.shape_inference import infer_shapes
 from typing import Dict, List, Any, Tuple, Sequence
 from functools import reduce
@@ -312,7 +318,7 @@ def from_onnx(model: ModelProto) -> backend.GraphHandler:
             raise Exception('Unsupported operator "{}"'.format(node.op_type))
 
 
-def to_onnx(graph: backend.GraphHandler):
+def to_onnx(graph: backend.GraphHandler, name: str) -> ModelProto:
     class Context:
         # saves object names, including tensors and operators
         names: Dict[Any, str] = dict()
@@ -370,6 +376,22 @@ def to_onnx(graph: backend.GraphHandler):
         def push_node(self, node: NodeProto) -> None:
             self.nodes.append(node)
 
+        def build(self, name: str) -> ModelProto:
+            print()
+            print(ctx.names)
+            print()
+            print(ctx.inputs)
+            print()
+            print(ctx.outputs)
+            print()
+            print(ctx.nodes)
+
+            return make_model(
+                make_graph(
+                    self.nodes, name, self.inputs, self.outputs, self.initializers
+                )
+            )
+
     # 拓扑排序
     if not graph.topo_sort():
         raise Exception("Sorting fails")
@@ -435,14 +457,7 @@ def to_onnx(graph: backend.GraphHandler):
         else:
             raise Exception("Unsupported OpType {}".format(ty.name))
 
-    print()
-    print(ctx.names)
-    print()
-    print(ctx.inputs)
-    print()
-    print(ctx.outputs)
-    print()
-    print(ctx.nodes)
+    return ctx.build(name)
 
 
 def parse_onnx(model: ModelProto):
diff --git a/pyinfinitensor/tests/test_onnx.py b/pyinfinitensor/tests/test_onnx.py
index a28dd5b9..d512b504 100644
--- a/pyinfinitensor/tests/test_onnx.py
+++ b/pyinfinitensor/tests/test_onnx.py
@@ -309,7 +309,7 @@ class TestStringMethods(unittest.TestCase):
         handler.add(abc, d, abcd)
         handler.add(abcd, e, abcde)
 
-        to_onnx(handler)
+        to_onnx(handler, "add")
 
 
 if __name__ == "__main__":