From 66b824237575ff12b05ef8d1f6f8d43ffe33e061 Mon Sep 17 00:00:00 2001 From: geohotstan <135171913+geohotstan@users.noreply.github.com> Date: Thu, 5 Dec 2024 23:31:26 +0800 Subject: [PATCH] Simple onnx.py clean ups (#8054) * start * simplify ops * why did this not work before * will split buffer parse to separate pr * flip the error order * only this much for now * to_python_const clean up * minimize diff * move tensor_methods into onnx.py * improve some type signatures --------- Co-authored-by: chenyu --- extra/onnx.py | 195 ++++++++++++++++++---------------------------- extra/onnx_ops.py | 20 +++-- 2 files changed, 89 insertions(+), 126 deletions(-) diff --git a/extra/onnx.py b/extra/onnx.py index 9e9ab1aa27..0b4493db2b 100644 --- a/extra/onnx.py +++ b/extra/onnx.py @@ -1,12 +1,10 @@ from __future__ import annotations -from typing import List, Dict, Union -import importlib -from functools import lru_cache +from typing import List, Dict, Union, Callable, Any +import importlib, functools import numpy as np -from tinygrad import Tensor, dtypes, Device -from tinygrad.tensor import _to_np_dtype -from tinygrad.helpers import getenv, DEBUG, CI, OSX -from tinygrad.dtype import ConstType, DType +from tinygrad import Tensor, dtypes +from tinygrad.helpers import getenv, DEBUG +from tinygrad.dtype import DType, ConstType from tinygrad.device import is_dtype_supported from onnx import AttributeProto, ModelProto, TensorProto, TypeProto try: @@ -17,32 +15,56 @@ except ImportError: def tensor_dtype_to_np_dtype(tensor_dtype:int) -> np.dtype: return TENSOR_TYPE_TO_NP_TYPE[tensor_dtype] cache_misses = 0 -@lru_cache(None) -def _cached_to_python_const(t:Tensor, tobytes): return t.data().tobytes() if tobytes else t.tolist() +@functools.lru_cache(None) +def _cached_to_python_const(t:Tensor): + if t.dtype is dtypes.uint8: return t.data().tobytes() + if 0 in t.shape: return [] + return t.tolist() # Tensor -> python value cache for parameters -def to_python_const(t, tobytes=False) -> Union[List[ConstType], List[bytes], Union[ConstType, bytes]]: +def to_python_const(t) -> Union[List[ConstType], List[bytes], Union[ConstType, bytes]]: if not isinstance(t, Tensor): return t global cache_misses - ret = _cached_to_python_const(t, tobytes) + ret = _cached_to_python_const(t) if (info := _cached_to_python_const.cache_info()).misses > cache_misses and DEBUG >= 3: - print(f"Cache miss for {t}, {tobytes=}") + print(f"Cache miss for {t}") cache_misses = info.misses return ret -# src: onnx/mapping.py https://onnx.ai/onnx/api/mapping.html#l-mod-onnx-mapping -# not supported: STRING = 8 COMPLEX64 = 14, COMPLEX128 = 15, UINT4 = 21, INT4 = 22 -# TODO: use dtypes.float16 for FLOAT16 -DTYPE_MAP: Dict[TensorProto.DataType, DType] = { - TensorProto.FLOAT:dtypes.float, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8, TensorProto.UINT16:dtypes.uint16, - TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64, TensorProto.BOOL:dtypes.bool, - TensorProto.FLOAT16:dtypes.float, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32, TensorProto.UINT64:dtypes.uint64, - TensorProto.BFLOAT16:dtypes.bfloat16, TensorProto.FLOAT8E4M3FN:dtypes.float, TensorProto.FLOAT8E4M3FNUZ:dtypes.float, - TensorProto.FLOAT8E5M2:dtypes.float, TensorProto.FLOAT8E5M2FNUZ:dtypes.float +# TODO: use real float16 +# src: onnx/mapping.py +DTYPE_MAP: Dict[TensorProto.DataType | int, DType] = { + TensorProto.FLOAT:dtypes.float32, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8, + TensorProto.UINT16:dtypes.uint16, TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64, + TensorProto.BOOL:dtypes.bool, TensorProto.FLOAT16:dtypes.float32, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32, + TensorProto.UINT64:dtypes.uint64, TensorProto.BFLOAT16:dtypes.bfloat16, TensorProto.FLOAT8E4M3FN:dtypes.float, + TensorProto.FLOAT8E4M3FNUZ:dtypes.float, TensorProto.FLOAT8E5M2:dtypes.float, TensorProto.FLOAT8E5M2FNUZ:dtypes.float } +def dtype_parse(onnx_dtype: TensorProto.DataType | int) -> DType: + if onnx_dtype not in DTYPE_MAP: raise NotImplementedError(f"onnx dtype {TensorProto.DataType.Name(onnx_dtype)} is not supported") + return DTYPE_MAP[onnx_dtype] if is_dtype_supported(DTYPE_MAP[onnx_dtype]) else dtypes.float + +# src: onnx/onnx_ml_pb2.pyi +ATTRIBUTE_MAP: Dict[AttributeProto.AttributeType, Callable[[AttributeProto], Any]] = { + AttributeProto.FLOAT: lambda a: float(a.f), AttributeProto.INT: lambda a: int(a.i), + AttributeProto.STRING: lambda a: a.s.decode("utf-8"), AttributeProto.TENSOR: lambda a: buffer_parse(a.t), + AttributeProto.FLOATS: lambda a: tuple(float(x) for x in a.floats), AttributeProto.INTS: lambda a: tuple(int(x) for x in a.ints), + AttributeProto.STRINGS: lambda a: tuple(x.decode("utf-8") for x in a.strings) +} +def attribute_parse(onnx_attribute: AttributeProto): + if onnx_attribute.type not in ATTRIBUTE_MAP: + raise NotImplementedError(f"attribute with type {AttributeProto.AttributeType.Name(onnx_attribute.type)} is not supported") + return ATTRIBUTE_MAP[onnx_attribute.type](onnx_attribute) + +def buffer_parse(inp: TensorProto) -> Tensor: + if dat := list(inp.float_data) or list(inp.int32_data) or list(inp.int64_data): + return Tensor(dat, dtype=dtype_parse(inp.data_type), requires_grad=False).reshape(tuple(inp.dims)) + if len(inp.raw_data) > 0: + return Tensor(np.frombuffer(inp.raw_data, dtype=tensor_dtype_to_np_dtype(inp.data_type)).copy().reshape(tuple(inp.dims)), + dtype=dtype_parse(inp.data_type), requires_grad=False) + raise NotImplementedError(f"buffer with data type {TensorProto.DataType.Name(inp.data_type)} is not supported") onnx_ops = importlib.import_module('extra.onnx_ops') - ONNXLIMIT = getenv("ONNXLIMIT", -1) def get_run_onnx(onnx_model: ModelProto): @@ -66,54 +88,30 @@ def get_run_onnx(onnx_model: ModelProto): elif attr == 'sparse_tensor_type': raise NotImplementedError(f"sparse_tensor_type is not implemented: {type_proto}") else: raise AttributeError(f"unknown attr: {attr}, {type_proto}") - def buffer_parse(inp: TensorProto) -> Tensor: - if inp.data_type not in DTYPE_MAP: - raise NotImplementedError(f"data type not supported {inp.name} {inp.dims} {inp.data_type}") - dtype = DTYPE_MAP[inp.data_type] if is_dtype_supported(DTYPE_MAP[inp.data_type]) else dtypes.float32 - if dat := list(inp.float_data) or list(inp.int32_data) or list(inp.int64_data): - return Tensor(dat, dtype=dtype, requires_grad=False).reshape(tuple(inp.dims)) - if len(inp.raw_data) > 0: - data = np.frombuffer(inp.raw_data, dtype=tensor_dtype_to_np_dtype(inp.data_type)).astype(_to_np_dtype(dtype)).copy() - return Tensor(data.reshape(tuple(inp.dims)), requires_grad=False) - return Tensor(None, requires_grad=False) - - def attribute_parse(a: AttributeProto) -> float | int | str | Tensor | tuple[float] | tuple[int]: - # TODO: this is not complete, see onnx/onnx_ml_pb2.pyi for a complete list - if a.type == AttributeProto.FLOAT: return float(a.f) - elif a.type == AttributeProto.INT: return int(a.i) - elif a.type == AttributeProto.STRING: return a.s.decode("utf-8") - elif a.type == AttributeProto.TENSOR: return buffer_parse(a.t) # TENSOR - elif a.type == AttributeProto.FLOATS: return tuple(float(x) for x in a.floats) - elif a.type == AttributeProto.INTS: return tuple(int(x) for x in a.ints) - elif a.type == AttributeProto.STRINGS: return tuple(x.decode("utf-8") for x in a.strings) - elif a.type == AttributeProto.GRAPH: raise NotImplementedError(f"graph not implemented: {a.g}\n likely an OP requiring control flow") - else: raise RuntimeError(f"can't parse {a.type} {a}") - - tensors: Dict[str, Tensor] = {} - - # get weights and biases - for inp in onnx_model.graph.initializer: - tensors[inp.name] = buffer_parse(inp) - - # preparse the attributes - attribute_dict = {} - domain = "" - for num,n in enumerate(onnx_model.graph.node): - attribute_dict[num] = {x.name:attribute_parse(x) for x in n.attribute} - if n.domain: domain = n.domain + # initialization data + model_parameters = {inp.name:buffer_parse(inp) for inp in onnx_model.graph.initializer} + model_attributes = {num:{x.name:attribute_parse(x) for x in n.attribute} for num,n in enumerate(onnx_model.graph.node)} + # model specs + is_onnx_preview_training = any(n.HasField("domain") and n.domain == "ai.onnx.preview.training" for n in onnx_model.graph.node) onnx_model_version = onnx_model.opset_import[0].version + # mapping from onnx ops to tensor.py ops + tensor_methods = { + op:op.lower() for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan", "Asin", "Acos", "Atan", + "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsInf", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh", "Tanh", + "Softsign", "Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Xor", "Round", "Erf") + } + def run_onnx(inputs={}, debug=0): debug = getenv("DEBUGONNX") or debug input_tensors: Dict[str,Tensor|List[Tensor]] = {} intermediate_tensors: Dict[str,Tensor] = {} - output_tensor_names = [x.name for x in onnx_model.graph.output] # get inputs for model_input in onnx_model.graph.input: name = model_input.name - if name in tensors: continue + if name in model_parameters: continue shape = type_parse(model_input.type) if name in inputs: if isinstance(inputs[name], Tensor): @@ -121,7 +119,7 @@ def get_run_onnx(onnx_model: ModelProto): elif isinstance(inputs[name], list): input_tensors[name] = [Tensor(i, requires_grad=False) for i in inputs[name]] # TODO: this is just to make training tests pass, need a principled way to handle training vs non-training - elif domain == "ai.onnx.preview.training": + elif is_onnx_preview_training: input_tensors[name] = Tensor(inputs[name], requires_grad=True) else: input_tensors[name] = Tensor(inputs[name], requires_grad=False) @@ -133,62 +131,27 @@ def get_run_onnx(onnx_model: ModelProto): raise RuntimeError(f"no data for {name} with shape {shape}") def fetch_tensor(x: str): - if x in tensors: return tensors[x] + 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 num,n in enumerate(onnx_model.graph.node): - inp: List[Tensor] = [] - if debug >= 3: print("inputs:") - for x in n.input: - t = fetch_tensor(x) - if debug >= 3: print(f"\t{x} - {t}") - inp.append(t) - opt: Dict = attribute_dict[num] - if debug >= 1: print(f"{num}: op {n.op_type} shape {[x.shape if isinstance(x, Tensor) else x for x in inp]} opt {opt}") + inp = [fetch_tensor(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}") + if debug >= 3: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {t}" for i,(x,t) in enumerate(zip(n.input, inp)))) + + if n.op_type in tensor_methods: + ret = getattr(Tensor, tensor_methods[n.op_type])(*inp, **opt) # NOTE some ops live here because they require access to some local variables - # have to use n.output for cases when num_outputs is absent - if n.op_type in onnx_ops.tensor_methods: - ret = getattr(Tensor, n.op_type.lower())(*inp, **opt) elif n.op_type == "Split": - axis = opt.get("axis", 0) - split = None if len(inp) == 1 else to_python_const(inp[1]) - if split is None: - split = [inp[0].shape[axis] // len(n.output)] * len(n.output) - for i in range(inp[0].shape[axis] % len(n.output)): - split[i] += 1 - i, ret = 0, [] - arg = [None] * inp[0].ndim - for s in split: - arg[axis] = (i,i+s) - ret.append(inp[0].shrink(arg=tuple(arg))) - i = i+s - ret = tuple(ret) - - # need to check onnx_model_version - elif n.op_type == "Slice": - if onnx_model_version < 10: - axes, ends, starts, steps = list(opt.get("axes", range(inp[0].ndim))), list(opt["ends"]), list(opt["starts"]), [1]*inp[0].ndim - else: - starts, ends = inp[1:3] - axes = list(range(inp[0].ndim)) if len(inp) <= 3 else to_python_const(inp[3].cast(dtypes.int32)) - steps = inp[4].cast(dtypes.int32).tolist() if len(inp) > 4 else [1]*inp[0].ndim - starts, ends = to_python_const(starts), to_python_const(ends) - arg = [(0,x,1) for x in inp[0].shape] - for i, axis in enumerate(axes): - axis = int(axis) + inp[0].ndim if axis < 0 else int(axis) - if starts[i] < 0: starts[i] += inp[0].shape[axis] - if ends[i] < 0: ends[i] += inp[0].shape[axis] - starts[i], ends[i] = max(0, min(starts[i], inp[0].shape[axis])), max(0, min(ends[i], inp[0].shape[axis])) - if starts[i] > ends[i] and steps[i] >= 0: steps[i] = -steps[i] - arg[axis] = (starts[i], ends[i], steps[i]) - new_shape = tuple((s, e) if st > 0 else (e+1, s+1) for s, e, st in arg) - if any(s==e for s,e in new_shape): ret = inp[0].shrink(new_shape) - else: ret = inp[0][tuple([slice(s,e,st) for s,e,st in arg])] - - # need to call backward on intermediate_tensors + axis, n_outputs = opt.get('axis', 0), opt.get('num_outputs') or len(n.output) + sz = inp[0].shape[axis] + sizes = to_python_const(inp[1]) if len(inp) == 2 else [sz // n_outputs + (1 if i < sz % n_outputs else 0) for i in range(n_outputs)] + ret = inp[0].split(sizes, axis) 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"] @@ -209,16 +172,12 @@ def get_run_onnx(onnx_model: ModelProto): print("UNSUPPORTED", n.op_type, n.input, n.output) raise NotImplementedError(f"op_type {n.op_type} not supported") + # finalization after running the op if not isinstance(ret, tuple): ret = (ret, ) - assert len(n.output) <= len(ret), f"expected output size must be less than {len(ret)}, it's {n.output}" - if debug >= 2: print([x.shape if isinstance(x, Tensor) else None for x in ret]) - if debug >= 2: print("outputs:") - for i in range(len(n.output)): - if debug >= 2: print(f"\t{n.output[i]} - {ret[i]}") - intermediate_tensors[n.output[i]] = ret[i] - if num == ONNXLIMIT: - output_tensor_names = n.output - break + 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] + if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{n.output[i]} - {ret[i]}" for i in range(len(n.output)))) - return {outp:intermediate_tensors[outp] for outp in output_tensor_names} + 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} return run_onnx diff --git a/extra/onnx_ops.py b/extra/onnx_ops.py index 06e4a7ae4b..7c80b99173 100644 --- a/extra/onnx_ops.py +++ b/extra/onnx_ops.py @@ -3,13 +3,9 @@ from typing import Union, Tuple, Optional, List, Any, cast from tinygrad.tensor import Tensor, _broadcast_shape from tinygrad.dtype import ImageDType, dtypes from tinygrad.helpers import prod, flatten -from extra.onnx import DTYPE_MAP, to_python_const +from extra.onnx import dtype_parse, to_python_const import numpy as np -tensor_methods = {"Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan", "Asin", "Acos", "Atan","Relu", - "Sigmoid", "MatMul", "Floor", "Ceil", "IsInf", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh", "Tanh", "Softsign", - "Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Xor", "Round", "Erf"} - # **************** Free Ops **************** def Identity(x: Tensor): return x @@ -26,7 +22,7 @@ def Min(*data_0): return functools.reduce(Tensor.minimum, data_0) def Sum(*data_0): return functools.reduce(Tensor.add, data_0) def Mean(*data_0): return Sum(*data_0) / len(data_0) # NOTE: does not support saturate -def Cast(x: Tensor, to: int, saturate=1): return x.cast(DTYPE_MAP[to]) +def Cast(x: Tensor, to: int, saturate=1): return x.cast(dtype_parse(to)) def CastLike(x: Tensor, target_type: Tensor, saturate=1): return x.cast(target_type.dtype) # **************** Simple Ops **************** @@ -91,6 +87,14 @@ def Trilu(x: Tensor, k: Union[Tensor, int]=0, upper=1): k = to_python_const(k) if isinstance(k, Tensor) else 0 # onnx passes k as a tensor int64 with one element, default is 0 return x.triu(k) if upper else x.tril(k) +def Slice(data: Tensor, starts:Tensor, ends:Tensor, axes:Optional[Tensor]=None, steps:Optional[Tensor]=None): + if axes is None: axes = list(range(data.ndim)) + if steps is None: steps = [1] * data.ndim + starts, ends, axes, steps = (to_python_const(x) for x in (starts, ends, axes, steps)) + slices = [slice(0,x,1) for x in data.shape] + for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i]) + return data[tuple(slices)] + def Squeeze(data: Tensor, axes): if isinstance(axes, Tensor): axes = to_python_const(axes) axes = [data._resolve_dim(x) for x in axes] @@ -405,7 +409,7 @@ def Compress(inp: Tensor, condition: Tensor, axis=None): return inp[tuple(con if i == axis else slice(None) for i in range(inp.ndim))] def EyeLike(x: Tensor, dtype=None, k=0): - ret = Tensor.eye(cast(int, min(x.shape)), dtype=DTYPE_MAP[dtype] if dtype else x.dtype) + ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype) if dtype else x.dtype) return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.size(0)-k) for d in x.shape)) def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode) @@ -424,7 +428,7 @@ def DequantizeLinear(x: Tensor, x_scale: Tensor, x_zero_point: Union[Tensor, int def ImageDecoder(encoded_stream: Tensor, pixel_format="RGB"): try: import PIL.Image except ImportError as e: raise ImportError("Pillow must be installed to use the reference implementation of the ImageDecoder operator") from e - img = PIL.Image.open(io.BytesIO(to_python_const(encoded_stream, True))) + img = PIL.Image.open(io.BytesIO(to_python_const(encoded_stream))) if pixel_format == "BGR": return Tensor(np.array(img))[:, :, ::-1] if pixel_format == "RGB": return Tensor(np.array(img)) if pixel_format == "Grayscale": return Tensor(np.array(img.convert("L"))).unsqueeze(-1) # (H, W) to (H, W, 1)