geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

unique names (#8213)

Browse Source
This commit is contained in:
geohotstan
2024-12-13 12:14:47 -05:00
committed by GitHub
parent 1824cbd72c
commit eebb3a1bb9
1 changed files with 10 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/onnx.py
+10 -17
View File
@@ -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