From eebb3a1bb9a7f6a85cf527e5b5cda85f9130896a Mon Sep 17 00:00:00 2001
From: geohotstan <135171913+geohotstan@users.noreply.github.com>
Date: Sat, 14 Dec 2024 01:14:47 +0800
Subject: [PATCH] unique names (#8213)

---
 extra/onnx.py | 27 ++++++++++-----------------
 1 file changed, 10 insertions(+), 17 deletions(-)

diff --git a/extra/onnx.py b/extra/onnx.py
index ad3a509067..521ea337be 100644
--- a/extra/onnx.py
+++ b/extra/onnx.py
@@ -69,7 +69,8 @@ ONNXLIMIT = getenv("ONNXLIMIT", -1)
 
 def get_run_onnx(onnx_model: ModelProto):
   # model initialization data
-  model_parameters = {inp.name:buffer_parse(inp) for inp in onnx_model.graph.initializer}
+  model_tensors = {inp.name:buffer_parse(inp) for inp in onnx_model.graph.initializer}
+  model_expected_inputs = {inp.name:inp for inp in onnx_model.graph.input if inp.name not in model_tensors}
   model_attributes = {num:{x.name:attribute_parse(x) for x in n.attribute} for num,n in enumerate(onnx_model.graph.node)}
 
   # model descriptions
@@ -114,21 +115,13 @@ def get_run_onnx(onnx_model: ModelProto):
 
   def run_onnx(inputs={}, debug=0):
     debug = getenv("DEBUGONNX") or debug
-    intermediate_tensors: Dict[str,Tensor] = {}
 
-    input_tensors: Dict[str, Tensor | List[Tensor]] = {}
-    for model_input in onnx_model.graph.input:
-      if model_input.name in inputs: input_tensors[model_input.name] = prepare_input(inputs[model_input.name], model_input)
-      elif model_input.name not in model_parameters: raise RuntimeError(f"Please provide input data for {model_input.name}")
-
-    def fetch_tensor(x: str):
-      if x in model_parameters: return model_parameters[x]
-      if x in intermediate_tensors: return intermediate_tensors[x]
-      if x != "": return input_tensors[x]
-      return None
+    for name, value_info in model_expected_inputs.items():
+      if name not in inputs: raise RuntimeError(f"Please provide input data for {name}")
+      model_tensors[name] = prepare_input(inputs[name], value_info)
 
     for num,n in enumerate(onnx_model.graph.node):
-      inp = [fetch_tensor(x) for x in n.input]
+      inp = [model_tensors.get(x) for x in n.input]
       opt = model_attributes[num]
 
       if debug >= 1: print(f"{num}: op \"{n.op_type}\" input shapes {[x.shape if isinstance(x, Tensor) else x for x in inp]} opt {opt}")
@@ -146,7 +139,7 @@ def get_run_onnx(onnx_model: ModelProto):
       elif n.op_type == "Gradient":
         assert len(opt["xs"]) == len(inp), f"len(opt['xs']):{len(opt['xs'])}, len(inp):{len(inp)} output and input has to match"
         y = opt["y"]
-        intermediate_tensors[y].backward()
+        model_tensors[y].backward()
         ret = tuple([t.grad for t in inp])
 
       # onnx_ops.py
@@ -166,9 +159,9 @@ def get_run_onnx(onnx_model: ModelProto):
       # finalization after running the op
       if not isinstance(ret, tuple): ret = (ret, )
       if len(n.output) > len(ret): raise RuntimeError(f"expected output size must be less than {len(ret)}, it's {n.output}")
-      for i in range(len(n.output)): intermediate_tensors[n.output[i]] = ret[i]
+      for i in range(len(n.output)): model_tensors[n.output[i]] = ret[i]
       if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{n.output[i]} - {ret[i]}" for i in range(len(n.output))))
 
-      if num == ONNXLIMIT: return {name:intermediate_tensors[name] for name in n.output}
-    return {x.name:intermediate_tensors[x.name] for x in onnx_model.graph.output}
+      if num == ONNXLIMIT: return {name:model_tensors[name] for name in n.output}
+    return {x.name:model_tensors[x.name] for x in onnx_model.graph.output}
   return run_onnx