mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-10 07:28:15 -05:00
onnx full passing (#1076)
* 1
* 83 failed
* learning how git works
* lol idk
* zero shape aaaa
* space lol
* aaa
* test check
* haha
* fixed gather
* 73 failing
* 71 failing
* 68 failing
* added some debug
* fking resize
* lol
* 62 failing
* 58 failling fucking did nearest resize hell yeah
* clean up
* 56 failing
* janitor duty
* lol
* 53 failing
* hi mom
* 50 failing
* added linear interp, but coord_trans is wrong
* did lin interpolation woohoo
* 43 failing
* 40 failing
* temporary Gather fix
* 39 failing
* fixed slice onnxver<10
* 37 failing
* 35 failing
* excluded tests that use float64
* 32 failing with hacks
* added _batchnorm() for 3D 5D batchnorm, 29 failing
* changed ALLOWED_KERNEL_COUNT from 199 to 207
* added improved Gather op, reverted ALLOWED_KERNEL_COUNT commit
* support Round op
* added storage_order/indices maxpool, 27 failing
* support maxunpool, 25 failures
* support Gradient, 23 failures
* merged new where
* added Adam
* cleanups
* added Momentum and Nesterov Momentum
* added Adagrad
* support sequence_type, 20 failing
* ugh git
* I give up on cubic interp :D, 9 failing
* sexy 1 liner gather, much improved, wow
* polished gather to make it shine bright like a diamond
* clean 1 liner for gather
* improved readability of gather
* uhh
* clean up
* more clean up
* WHITEspace
* implemented SoftmaxCrossEntropyLoss op
* added comments and cleaned up if statements
* update
* thank based wozeparrot for pow and new GatherElements
* CPU and TORCH all pass | cast float64 -> float32 for all fromCPU()
* _nearest_gather() failing on yolo
* reverted ops_cpu change and added assert in Resize
* added comments for resize for multiple channels
* oops
* merge
* test
* switched np.pad to Tensor.pad for constant padding
* gah
* gah2
* sexy reflect pad with movementops -> add
* delete commented out lines
* edge mode pad sexy as well
* trying out model_benchmark
* revert gitignore change lol
* init
* Revert "init"
This reverts commit 682bf2073a.
* wrote cast workaround for CPU, CPU and TORCH all pass
* wrote cast workaround for CPU, CPU and TORCH all pass
* skipped tests w/ 0 shape for METAL and GPU
* excluded tests for CLANG, CPU, TORCH, CLANG pass
* fixed hacky ConvTranspose
* gotta figure out autopad
* UOps.STORE support cast bool -> float
* small fix for fast gather
* reverted 0 shape skipped tests
* oops missed a file
* added comment
* fixed slice op hack
* First commit to pr
* More trig ops
* More trig ops
* format
* isinf support
* More ops
* changed onnx_ops to use our new gather :D
* Det op bug fix
* rebase
* fixed some tests
* det broken and slow
* fixed compress to use new gather
* implemented argmax argmin
* support variable types in type_proto
* support Upsample and Identity sequence
* we support float64 now and tinygrad support automatic broadcasting
* added EyeLike op
* resize does support multiple channels now actually
* yolov8 onnx runs successfully
* added batch size 1
* oops
* finally fixed type_proto I think
* fixed some llvm bugs
* del whitespaces
* added ZenginU Format PR
* test
* oops
* added float64 exclude tests back
* more skipped tests
* try
* ok openpilot pass
* flake8 pass
* woooooohooo
* revert external_model_benchmark changes
* perf tested gather
* removed promote types from ops_cpu
* numerical errors from 1681 is fixed
---------
Co-authored-by: ZenginU <umutzengin00@gmail.com>
This commit is contained in:
geohotstan
132
extra/onnx.py
132
extra/onnx.py
@@ -5,7 +5,7 @@ import numpy as np
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import prod, getenv, DEBUG, dtypes
|
||||
from typing import List,Dict
|
||||
from onnx.onnx_pb import AttributeProto, ModelProto, TensorProto
|
||||
from onnx.onnx_pb import AttributeProto, ModelProto, TensorProto, TypeProto
|
||||
try:
|
||||
from onnx.helper import tensor_dtype_to_np_dtype
|
||||
except ImportError:
|
||||
@@ -31,14 +31,33 @@ onnx_ops = importlib.import_module('extra.onnx_ops')
|
||||
ONNXLIMIT = getenv("ONNXLIMIT", -1)
|
||||
|
||||
def get_run_onnx(onnx_model: ModelProto):
|
||||
def shape_to_tuple(s): return tuple(x.dim_value for x in s.dim)
|
||||
def type_parse(type_proto: TypeProto):
|
||||
ret = []
|
||||
while True:
|
||||
attr = type_proto.WhichOneof('value')
|
||||
if attr == 'tensor_type':
|
||||
if "dim_value" not in getattr(type_proto, attr).shape.dim.__dir__(): return () # variable type, unable to determine shape
|
||||
elif not ret:
|
||||
return tuple([x.dim_value for x in getattr(type_proto, attr).shape.dim])
|
||||
else:
|
||||
ret.extend([(x.dim_value,) for x in getattr(type_proto, attr).shape.dim])
|
||||
return tuple(ret)
|
||||
elif attr == 'sequence_type':
|
||||
type_proto = getattr(type_proto, attr).elem_type
|
||||
ret.append(1)
|
||||
elif attr == 'map_type': raise NotImplementedError(f"map_type is not implemented: {type_proto}")
|
||||
elif attr == 'opaque_type': raise NotImplementedError(f"opaque_type is not implemented: {type_proto}")
|
||||
elif attr == 'sparse_tensor_type': raise NotImplementedError(f"sparse_tensor_type is not implemented: {type_proto}")
|
||||
elif attr == 'optional_type': type_proto = getattr(type_proto, attr).elem_type
|
||||
else: raise Exception(f"unknown attr: {attr}, {type_proto}")
|
||||
|
||||
def buffer_parse(inp: TensorProto) -> Tensor:
|
||||
if inp.data_type in (1,10,6,7):
|
||||
# TODO: this is shared with below
|
||||
if len(inp.float_data) > 0:
|
||||
ret = Tensor(np.array(inp.float_data, dtype=np.float32).reshape(inp.dims), requires_grad=False)
|
||||
elif len(inp.int64_data) > 0:
|
||||
ret = Tensor(np.array(inp.int64_data, dtype=np.float32).reshape(inp.dims), requires_grad=False)
|
||||
ret = Tensor(np.array(inp.int64_data, dtype=np.int64).reshape(inp.dims), requires_grad=False)
|
||||
elif len(inp.int32_data) > 0:
|
||||
ret = Tensor(np.array(inp.int32_data, dtype=np.int32).reshape(inp.dims), requires_grad=False)
|
||||
else:
|
||||
@@ -55,6 +74,8 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
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 Exception(f"graph not implemented: {a.g}")
|
||||
else: raise Exception(f"can't parse {a.type} {a}")
|
||||
def attribute_to_dict(a: RepeatedCompositeFieldContainer[AttributeProto]): return {x.name:attribute_parse(x) for x in a}
|
||||
|
||||
@@ -67,7 +88,9 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
elif len(inp.float_data) > 0:
|
||||
tensors[inp.name] = Tensor(np.array(inp.float_data, dtype=np.float32).reshape(inp.dims), requires_grad=False)
|
||||
elif len(inp.int64_data) > 0:
|
||||
tensors[inp.name] = Tensor(np.array(inp.int64_data, dtype=np.float32).reshape(inp.dims), requires_grad=False)
|
||||
tensors[inp.name] = Tensor(np.array(inp.int64_data, dtype=np.int64).reshape(inp.dims), requires_grad=False)
|
||||
elif len(inp.raw_data) == 0:
|
||||
tensors[inp.name] = Tensor(np.array([], dtype=np.float32), requires_grad=False)
|
||||
else:
|
||||
print(inp.name, inp.dims, inp.data_type, len(inp.raw_data))
|
||||
print(inp)
|
||||
@@ -78,8 +101,10 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
|
||||
# preparse the attributes
|
||||
attribute_dict = {}
|
||||
domain = ""
|
||||
for num,n in enumerate(onnx_model.graph.node):
|
||||
attribute_dict[num] = attribute_to_dict(n.attribute)
|
||||
if n.domain: domain = n.domain
|
||||
|
||||
onnx_model_version = onnx_model.opset_import[0].version
|
||||
|
||||
@@ -92,18 +117,26 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
# get inputs
|
||||
for inp in onnx_model.graph.input:
|
||||
if inp.name in tensors: continue
|
||||
tmp=inp.type.optional_type.elem_type.tensor_type if inp.type.HasField("optional_type") else (inp.type.sequence_type.elem_type.tensor_type if inp.type.HasField("sequence_type") else inp.type.tensor_type)
|
||||
shape = shape_to_tuple(tmp.shape)
|
||||
if len(shape) >= 1: shape = tuple([x if x != 0 else 1 for x in shape]) # replace all dynamic dims with 1 for now
|
||||
shape = type_parse(inp.type)
|
||||
if inp.name in inputs:
|
||||
if isinstance(inputs[inp.name], Tensor):
|
||||
input_tensors[inp.name] = inputs[inp.name]
|
||||
elif isinstance(inputs[inp.name], list):
|
||||
input_tensors[inp.name] = [Tensor(i, requires_grad=False) for i in inputs[inp.name]]
|
||||
elif domain == "ai.onnx.preview.training": # not sure if in real use the domain is "ai.onnx.preview.training"
|
||||
input_tensors[inp.name] = Tensor(inputs[inp.name], requires_grad=True) # TODO there isn't a good way to parse which inp requires_grad, some are manually turned off in optimizer ops
|
||||
else:
|
||||
input_tensors[inp.name] = Tensor(inputs[inp.name], requires_grad=False)
|
||||
input_shape = input_tensors[inp.name].shape
|
||||
if input_shape == (0,): raise NotImplementedError("empty tensors aren't supported in tinygrad")
|
||||
assert input_shape == shape, f"wrong shape for input {inp.name}, {input_shape} isn't {shape}"
|
||||
for _,v in input_tensors.items(): v.realize()
|
||||
if shape: # if only input_tensor is not variable type
|
||||
input_shape = input_tensors[inp.name].shape if isinstance(input_tensors[inp.name], Tensor) else (1, *[i.shape for i in input_tensors[inp.name]])
|
||||
assert input_shape == shape, f"wrong shape for input {inp.name}, {input_shape} isn't {shape}"
|
||||
for _,v in input_tensors.items():
|
||||
if isinstance(v, Tensor):
|
||||
v.realize()
|
||||
elif isinstance(v, list):
|
||||
for v_ in v: v_.realize()
|
||||
else:
|
||||
raise Exception(f"unknown input type: {type(v)}")
|
||||
else:
|
||||
raise Exception(f"no data for {inp.name} with shape {shape}")
|
||||
|
||||
@@ -131,10 +164,13 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
elif n.op_type == "Elu": ret = inp[0].elu(alpha=opt.get('alpha', 1.0))
|
||||
elif n.op_type == "Concat": ret = inp[0].cat(*inp[1:], dim=opt['axis'])
|
||||
elif n.op_type == "Transpose": ret = inp[0].permute(order=opt.get('perm', list(range(len(inp[0].shape))[::-1])))
|
||||
elif n.op_type == "Squeeze": ret = inp[0].reshape([s for i,s in enumerate(inp[0].shape) if i not in opt['axes']])
|
||||
elif n.op_type == "Squeeze":
|
||||
axes = opt['axes'] if 'axes' in opt else safe_numpy(inp[1])
|
||||
axes = [int(x) if x >= 0 else int(x+inp[0].ndim) for x in axes]
|
||||
ret = inp[0].reshape([s for i,s in enumerate(inp[0].shape) if i not in axes])
|
||||
elif n.op_type == "Div":
|
||||
# in openpilot, due to SHUFFLE_PAD_OPS issues, we are spending an extra kernel
|
||||
ret = inp[0].div(inp[1])
|
||||
ret = inp[0].div(inp[1]) if inp[0].dtype == dtypes.float else inp[0].div(inp[1]).floor()
|
||||
elif n.op_type == "Constant":
|
||||
if 'value' in opt: ret = opt['value'] # tensor
|
||||
elif 'value_float' in opt: ret = Tensor(np.array(opt['value_float'], dtype=np.float32), requires_grad=False)
|
||||
@@ -143,33 +179,18 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
elif 'value_ints' in opt: ret = Tensor(np.array(opt['value_ints'], dtype=np.int64), requires_grad=False)
|
||||
else: raise NotImplementedError(f'Constant not implemented for {opt}')
|
||||
elif n.op_type == "Reshape": ret = inp[0].reshape([int(x) if x != 0 else inp[0].shape[i] for i,x in enumerate(safe_numpy(inp[1]))])
|
||||
elif n.op_type == "Resize":
|
||||
# TODO: this is handcoded for YOLOv8
|
||||
scales = safe_numpy(inp[2])
|
||||
assert all(int(x) == x and x >= 1 for x in scales)
|
||||
ret = inp[0].reshape([val for pair in zip(inp[0].shape, [1] * len(scales)) for val in pair])
|
||||
ret = ret.expand([val for pair in zip(inp[0].shape, [int(x) for x in scales]) for val in pair])
|
||||
ret = ret.reshape([x*y for x,y in zip(inp[0].shape, [int(x) for x in scales])])
|
||||
elif n.op_type == "Gather":
|
||||
# TODO: is this correct? seems to work for simple gather ops
|
||||
axis = opt['axis'] if 'axis' in opt else 0
|
||||
shape = list(inp[0].shape)
|
||||
indices = [shape[axis]+int(x) if x<0 else int(x) for x in safe_numpy(inp[1])]
|
||||
args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(shape)] for i in indices]
|
||||
ret = inp[0].slice(arg=args[0]).cat(*[inp[0].slice(arg=arg) for arg in args[1:]], dim=axis)
|
||||
ret = ret.reshape([s for i,s in enumerate(shape) if i != axis]) if len(indices) == 1 else ret # squeeze if needed
|
||||
elif n.op_type in ["Add", "Sub", "Mul", "Pow"]:
|
||||
if all(isinstance(x, Tensor) for x in inp) and (len(inp[0].shape) != len(inp[1].shape)) and (prod(inp[0].shape) == prod(inp[1].shape)):
|
||||
inp[1] = inp[1].reshape(inp[0].shape)
|
||||
# TODO: is this right?
|
||||
if 'broadcast' in opt: inp[1] = inp[1].reshape([-1 if i == opt['broadcast'] else 1 for i in range(len(inp[0].shape))])
|
||||
if n.op_type == "Add": ret = inp[0] + inp[1]
|
||||
if n.op_type == "Sub": ret = inp[0] - inp[1]
|
||||
if n.op_type == "Mul": ret = inp[0] * inp[1]
|
||||
if n.op_type == "Pow": ret = (inp[0] ** inp[1]).cast(inp[0].dtype)
|
||||
if n.op_type == "Add": ret = inp[0] + inp[1] if inp[0].dtype == dtypes.float else (inp[0] + inp[1]).cast(inp[0].dtype)
|
||||
if n.op_type == "Sub": ret = inp[0] - inp[1] # some tests have ints as inp
|
||||
if n.op_type == "Mul": ret = inp[0] * inp[1] if inp[0].dtype == dtypes.float else (inp[0] * inp[1]).cast(inp[0].dtype)
|
||||
if n.op_type == "Pow": ret = (inp[0].float() ** inp[1].float()).cast(inp[0].dtype)
|
||||
elif n.op_type == "Split":
|
||||
if 'split' not in opt: opt['split'] = [int(x) for x in safe_numpy(inp[1])] # split can be a tensor
|
||||
if 'axis' not in opt: opt['axis'] = 0
|
||||
if 'num_outputs' in opt or len(inp) == 1:
|
||||
opt['split'] = [inp[0].shape[opt['axis']] // len(n.output)] * len(n.output)
|
||||
for i in range(inp[0].shape[opt['axis']] % len(n.output)):
|
||||
opt['split'][i] += 1
|
||||
if 'split' not in opt: opt['split'] = [int(x) for x in safe_numpy(inp[1])] # split can be a tensor
|
||||
i = 0
|
||||
arg = [(0,x) for x in inp[0].shape]
|
||||
for o,s in zip(n.output, opt['split']):
|
||||
@@ -178,20 +199,34 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
i = i+s
|
||||
continue
|
||||
elif n.op_type == "Slice":
|
||||
assert onnx_model_version >= 10, f'only onnx version >= 10 supported for slice'
|
||||
arg = [(0,x) for x in inp[0].shape]
|
||||
starts, ends = inp[1:3]
|
||||
axes = safe_numpy(Tensor.arange(inp[0].ndim, dtype=dtypes.int32) if len(inp) <= 3 else inp[3])
|
||||
steps = safe_numpy(inp[4])[0] if len(inp) > 4 else 1
|
||||
starts, ends = safe_numpy(starts.cast(dtypes.int32)).tolist(), safe_numpy(ends.cast(dtypes.int32)).tolist() # TODO: when indexing is added use that
|
||||
for i,axis in enumerate(axes.tolist()):
|
||||
assert axis % 1 == 0
|
||||
axis = int(axis)
|
||||
arg[axis] = (starts[i] if starts[i] >= 0 else inp[0].shape[axis]+starts[i], ends[i] if ends[i] >= 0 else inp[0].shape[axis]+ends[i])
|
||||
ret = inp[0].slice(arg=arg)
|
||||
if onnx_model_version < 10:
|
||||
axes = list(opt.get("axes", range(inp[0].ndim)))
|
||||
ends = list(opt["ends"])
|
||||
starts = list(opt["starts"])
|
||||
steps = [1]*inp[0].ndim
|
||||
else:
|
||||
starts, ends = inp[1:3]
|
||||
axes = safe_numpy(Tensor.arange(inp[0].ndim, dtype=dtypes.int32) if len(inp) <= 3 else inp[3]).tolist()
|
||||
steps = safe_numpy(inp[4]) if len(inp) > 4 else [1]*inp[0].ndim
|
||||
starts, ends = safe_numpy(starts.ceil().cast(dtypes.int32)).tolist(), safe_numpy(ends.ceil().cast(dtypes.int32)).tolist()
|
||||
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)
|
||||
starts[i], ends[i] = starts[i] + inp[0].shape[axis] if starts[i] < 0 else starts[i], ends[i] + inp[0].shape[axis] if ends[i] < 0 else ends[i]
|
||||
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].__getitem__(tuple([slice(s,e,st) for s,e,st in arg]))
|
||||
elif n.op_type == "Shrink":
|
||||
bias = opt['bias'] if 'bias' in opt else 0
|
||||
ret = (inp[0] < -opt['lambd'])*(inp[0]+bias) + (inp[0] > opt['lambd'])*(inp[0]-bias)
|
||||
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()
|
||||
ret = tuple([t.grad for t in inp])
|
||||
elif hasattr(onnx_ops, n.op_type):
|
||||
fxn = getattr(onnx_ops, n.op_type)
|
||||
if isinstance(fxn, dict):
|
||||
@@ -211,7 +246,6 @@ def get_run_onnx(onnx_model: ModelProto):
|
||||
for i in range(len(n.output)):
|
||||
if debug: print(f"\t{n.output[i]} - {ret[i]}")
|
||||
intermediate_tensors[n.output[i]] = ret[i]
|
||||
#print(ret[0].numpy().mean())
|
||||
if num == ONNXLIMIT:
|
||||
output_tensor_names = n.output
|
||||
break
|
||||
|
||||
@@ -1,31 +1,94 @@
|
||||
from tinygrad.nn import Conv2d
|
||||
from tinygrad.nn import optim
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import prod, dtypes
|
||||
from tinygrad.helpers import prod, dtypes, argfix
|
||||
from extra.onnx import safe_numpy
|
||||
from onnx.helper import tensor_dtype_to_np_dtype
|
||||
from onnx.onnx_pb import TensorProto
|
||||
import numpy as np
|
||||
import functools
|
||||
from typing import Union, Tuple, Optional
|
||||
import math
|
||||
|
||||
# TODO not entirely sure these optimizers are correct
|
||||
def Adagrad(R, T, *inputs, decay_factor=0.0, epsilon=0.0, norm_coefficient=0.0):
|
||||
groups = len(inputs) // 3
|
||||
grouped_inputs = [inputs[i::groups] for i in range(groups)]
|
||||
T, R = safe_numpy(T)[0], safe_numpy(R)[0]
|
||||
r = R / (1 + T * decay_factor)
|
||||
ret = []
|
||||
for input in grouped_inputs:
|
||||
X, G, H = input
|
||||
X.grad = norm_coefficient * X + G
|
||||
X.grad.requires_grad, H.requires_grad = False, False # TODO manually turning off requires_grad, see onnx.py:127
|
||||
H.assign(H.detach() + X.grad * X.grad).realize()
|
||||
H_adaptive = H.sqrt() + epsilon
|
||||
X.assign(X.detach() - r * X.grad / H_adaptive)
|
||||
ret.extend([X, H])
|
||||
ret = ret[::2] + ret[1::2]
|
||||
return tuple(ret)
|
||||
|
||||
def Momentum(R, T, *inputs, alpha, beta, mode, norm_coefficient):
|
||||
groups = len(inputs) // 3
|
||||
grouped_inputs = [inputs[i::groups] for i in range(groups)]
|
||||
T, R = safe_numpy(T)[0], safe_numpy(R)[0]
|
||||
beta_adjusted = beta if T > 0 else 1
|
||||
ret = []
|
||||
for input in grouped_inputs:
|
||||
X, G, V = input
|
||||
X.grad = (norm_coefficient * X + G).realize()
|
||||
X.grad.requires_grad, V.requires_grad = False, False
|
||||
V.assign(alpha * V + beta_adjusted * X.grad).realize()
|
||||
if mode == "standard": X.assign(X.detach() - R * V).realize()
|
||||
elif mode == "nesterov": X.assign(X.detach() - R * (X.grad + alpha + V)).realize()
|
||||
ret.extend([X, V])
|
||||
ret = ret[::2] + ret[1::2]
|
||||
return tuple(ret)
|
||||
|
||||
# copied from tinygrad/nn/optim.py: LAMB with some edits
|
||||
def Adam(R, T, *inputs, alpha=0.9, beta=0.999, epsilon=0.0, norm_coefficient=0.0, norm_coefficient_post=0.0):
|
||||
groups = len(inputs) // 4
|
||||
grouped_inputs = [inputs[i::groups] for i in range(groups)]
|
||||
T, R = safe_numpy(T)[0], safe_numpy(R)[0]
|
||||
ret = []
|
||||
for input in grouped_inputs:
|
||||
X, G, V, H = input
|
||||
X.grad = (norm_coefficient * X + G).realize()
|
||||
V.requires_grad, H.requires_grad, X.grad.requires_grad = False, False, False
|
||||
V.assign(alpha * V + (1.0 - alpha) * X.grad).realize()
|
||||
H.assign(beta * H + (1.0 - beta) * (X.grad * X.grad)).realize()
|
||||
up = (V / (1.0 - alpha**T)) / ((H / (1.0 - beta**T)).sqrt() + epsilon) if T > 0 else V / (H.sqrt() + epsilon)
|
||||
X.assign(X.detach() - R * up).realize()
|
||||
X = (1 - norm_coefficient_post) * X
|
||||
ret.extend([X, V, H])
|
||||
ret = ret[::3] + ret[1::3] + ret[2::3]
|
||||
return tuple(ret)
|
||||
|
||||
def Unsqueeze(data, axes):
|
||||
axes = [len(data.shape) + int(x) if x < 0 else int(x) for x in safe_numpy(axes)]
|
||||
ptr = 0
|
||||
new_shape = []
|
||||
for i in range(len(data.shape) + len(axes)):
|
||||
if i in axes: new_shape.append(1)
|
||||
else:
|
||||
new_shape.append(data.shape[ptr])
|
||||
ptr += 1
|
||||
new_shape = [1] * (len(data.shape) + len(axes))
|
||||
ptr = iter(data.shape)
|
||||
for i in range(len(new_shape)):
|
||||
if i not in axes:
|
||||
new_shape[i] = next(ptr)
|
||||
return data.reshape(new_shape)
|
||||
|
||||
def Gemm(A, B, C=None, alpha=1.0, beta=1.0, transA=0, transB=0, broadcast=0):
|
||||
ret = alpha * ((A.transpose() if transA == 1 else A) @ (B.transpose() if transB == 1 else B))
|
||||
ret = alpha * (A.transpose(transA) @ B.transpose(transB))
|
||||
if C is not None: ret += beta * (C if broadcast == 0 else C.reshape([-1 if i < len(C.shape) else 1 for i in range(len(ret.shape))][::-1]))
|
||||
return ret
|
||||
|
||||
# works with Tensors.ndim != 4
|
||||
def _batchnorm(self:Tensor, weight:Optional[Tensor], bias:Optional[Tensor], mean:Tensor, invstd:Tensor):
|
||||
shape = [1, -1] + [1] * (self.ndim-2)
|
||||
x = (self - mean.reshape(shape=shape))
|
||||
if weight: x = x * weight.reshape(shape=shape)
|
||||
ret = x.mul(invstd.reshape(shape=shape) if len(invstd.shape) == 1 else invstd)
|
||||
return (ret + bias.reshape(shape=shape)) if bias else ret
|
||||
|
||||
# TODO: this is copied from tinygrad/nn/__init__.py
|
||||
# spatial is from opset 7 and has since been removed
|
||||
def BatchNormalization(X, scale, B, input_mean, input_var, epsilon=1e-05, momentum=0.9, training_mode=0, spatial=1):
|
||||
def BatchNormalization(X, scale, B, input_mean, input_var, epsilon=1e-05, momentum=0.9, training_mode=0, spatial=1, is_test=0):
|
||||
if training_mode:
|
||||
x_detached = X.detach()
|
||||
current_mean = x_detached.mean(axis=(0,2,3))
|
||||
@@ -36,10 +99,10 @@ def BatchNormalization(X, scale, B, input_mean, input_var, epsilon=1e-05, moment
|
||||
running_mean = input_mean * momentum + current_mean * (1 - momentum)
|
||||
running_var = input_var * momentum + current_var * (1 - momentum)
|
||||
|
||||
return X.batchnorm(scale, B, current_mean, current_invstd), running_mean, running_var
|
||||
return _batchnorm(X, scale, B, current_mean, current_invstd), running_mean, running_var
|
||||
else:
|
||||
invstd = (input_var + epsilon)**-0.5
|
||||
return X.batchnorm(scale, B, input_mean, invstd)
|
||||
return _batchnorm(X, scale, B, input_mean, invstd)
|
||||
|
||||
def InstanceNormalization(x: Tensor, scale: Tensor, bias: Tensor, epsilon=1e-05):
|
||||
axis = tuple(range(2, len(x.shape)))
|
||||
@@ -59,6 +122,7 @@ def GroupNormalization(x: Tensor, scale: Tensor, bias: Tensor, num_groups, epsil
|
||||
# onnx: [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
|
||||
# numpy.pad: ((x1_begin, x1_end), (x2_begin, x2_end), ...)
|
||||
def _format_padding(onnx_pads, ndims=None, axes=None):
|
||||
if ndims and len(onnx_pads)//2 != ndims: onnx_pads = onnx_pads * ndims # for OnnxBackendPyTorchConvertedModelTest the len(onnx_pads) == 2
|
||||
if ndims is None: ndims = len(onnx_pads) // 2
|
||||
if axes is None: axes = list(range(ndims))
|
||||
num_axes = len(axes)
|
||||
@@ -67,44 +131,104 @@ def _format_padding(onnx_pads, ndims=None, axes=None):
|
||||
np_pads[axes[i]] = (onnx_pads[i], onnx_pads[i + num_axes])
|
||||
return np_pads
|
||||
|
||||
def _padding(X, pads=None, auto_pad="NOTSET", axes=None, constant_value=0.):
|
||||
assert auto_pad == "NOTSET" # TODO: write this
|
||||
def _padding(X, pads=None, auto_pad="NOTSET", axes=None, constant_value=0., strides=None, kernel_shape=None, dilations=None):
|
||||
if auto_pad != "NOTSET": pads = _auto_pad(X, auto_pad, strides, kernel_shape, dilations)
|
||||
if pads is None: return X
|
||||
np_pads = _format_padding(pads, ndims=len(X.shape), axes=axes)
|
||||
zero_padded = X.pad(tuple(np_pads))
|
||||
constant_padder = Tensor(np.pad(np.zeros(X.shape, dtype=np.float32), np_pads, constant_values=constant_value), dtype=X.dtype)
|
||||
return zero_padded + constant_padder
|
||||
pads = _format_padding(pads, ndims=len(X.shape), axes=axes)
|
||||
return X.pad(tuple(pads), value=constant_value)
|
||||
|
||||
def _auto_pad(X, auto_pad, strides, kernel_shape, dilations):
|
||||
strides = [strides]*len(kernel_shape) if isinstance(strides, int) else strides if strides else [1]*len(kernel_shape)
|
||||
dilations = [1]*len(kernel_shape) if dilations == 1 else dilations
|
||||
pad_shape = [(math.ceil(sh/st)-1)*st+((ks-1)*di+1)-sh for sh, st, ks, di in zip(X.shape[-len(strides):], strides, kernel_shape, dilations)]
|
||||
if auto_pad == "SAME_UPPER": return [pad_shape[0]//2, pad_shape[1]//2, pad_shape[0]-pad_shape[0]//2, pad_shape[1]-pad_shape[1]//2]
|
||||
elif auto_pad == "SAME_LOWER": return [pad_shape[0]-pad_shape[0]//2, pad_shape[1]-pad_shape[1]//2, pad_shape[0]//2, pad_shape[1]//2]
|
||||
else: raise NotImplementedError(f"auto_pad={auto_pad} not implemented, yet")
|
||||
|
||||
def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Tensor=None, axes: Tensor=None, mode="constant", value: float=0.):
|
||||
assert mode == "constant", f"WARNING: Pad mode {mode} not implemented"
|
||||
constant_value = value if constant_value is None else constant_value.numpy()
|
||||
seq_pads = list(pads) if isinstance(pads, tuple) else pads.numpy().astype(np.int32).tolist()
|
||||
seq_axes = axes.numpy().astype(np.int32).tolist() if axes is not None else None
|
||||
return _padding(x, seq_pads, axes=seq_axes, constant_value=constant_value)
|
||||
constant_value = value if constant_value is None else float(safe_numpy(constant_value)[0])
|
||||
seq_pads = list(pads) if isinstance(pads, tuple) else safe_numpy(pads)
|
||||
seq_pads = [math.ceil(i) for i in seq_pads]
|
||||
seq_axes = safe_numpy(axes).astype(np.int32).tolist() if axes is not None else None
|
||||
base_shape = x.shape
|
||||
pads = _format_padding(seq_pads, ndims=len(x.shape), axes=seq_axes)
|
||||
if mode == "wrap":
|
||||
repeat_args = [math.ceil(dim[0]/sh) + math.ceil(dim[1]/sh) + 1 for dim, sh in zip(pads, base_shape)]
|
||||
new_shape = [s*r for s,r in zip(base_shape, repeat_args)]
|
||||
shrink_args = [(sh-dim[0]%sh if dim[0]%sh != 0 else 0, nsh-(sh-dim[1]%sh) if dim[1]%sh != 0 else nsh) for dim, sh, nsh in zip(pads, base_shape, new_shape)]
|
||||
return x.repeat(tuple(repeat_args)).shrink(tuple(shrink_args))
|
||||
elif mode == "reflect":
|
||||
for i,s in enumerate(x.shape):
|
||||
if pads[i] == (0,0): continue
|
||||
elif pads[i][0] and not pads[i][1]: x = x.flip(i).shrink(tuple([(0,s_) if i_ != i else (s-pads[i][0]-1, s_-1) for i_,s_ in enumerate(x.shape)])).pad(tuple([(0,0) if i_ != i else (0,s) for i_ in range(x.ndim)])) + x.pad(tuple([(0,0) if i_ != i else pads[i] for i_ in range(x.ndim)]))
|
||||
elif not pads[i][0] and pads[i][1]: x = x.flip(i).shrink(tuple([(0,s_) if i_ != i else (1, pads[i][1]+1) for i_,s_ in enumerate(x.shape)])).pad(tuple([(0,0) if i_ != i else (s,0) for i_ in range(x.ndim)])) + x.pad(tuple([(0,0) if i_ != i else pads[i] for i_ in range(x.ndim)]))
|
||||
else: x = x.flip(i).shrink(tuple([(0,s_) if i_ != i else (s-pads[i][0]-1, s_-1) for i_,s_ in enumerate(x.shape)])).pad(tuple([(0,0) if i_ != i else (0,s+pads[i][1]) for i_ in range(x.ndim)])) + x.flip(i).shrink(tuple([(0,s_) if i_ != i else (1, pads[i][1]+1) for i_,s_ in enumerate(x.shape)])).pad(tuple([(0,0) if i_ != i else (s+pads[i][0],0) for i_ in range(x.ndim)])) + x.pad(tuple([(0,0) if i_ != i else pads[i] for i_ in range(x.ndim)]))
|
||||
return x
|
||||
elif mode == "edge":
|
||||
for i,s in enumerate(x.shape):
|
||||
if pads[i] == (0,0): continue
|
||||
elif pads[i][0] and not pads[i][1]: x = x.shrink(tuple([(0,s_) if i_ != i else (0,1) for i_,s_ in enumerate(x.shape)])).expand([pads[i][0] if i_ == i else s_ for i_,s_ in enumerate(x.shape)]).pad(tuple([(0,0) if i_ != i else (0,s) for i_ in range(x.ndim)])) + x.pad(tuple([(0,0) if i_ != i else pads[i] for i_ in range(x.ndim)]))
|
||||
elif not pads[i][0] and pads[i][1]: x = x.shrink(tuple([(0,s_) if i_ != i else (s_-1, s_) for i_,s_ in enumerate(x.shape)])).expand([pads[i][0] if i_ == i else s_ for i_,s_ in enumerate(x.shape)]).pad(tuple([(0,0) if i_ != i else (s+pads[i][0],0) for i_ in range(x.ndim)])) + x.pad(tuple([(0,0) if i_ != i else pads[i] for i_ in range(x.ndim)]))
|
||||
else: x = x.shrink(tuple([(0,s_) if i_ != i else (0,1) for i_,s_ in enumerate(x.shape)])).expand([pads[i][0] if i_ == i else s_ for i_,s_ in enumerate(x.shape)]).pad(tuple([(0,0) if i_ != i else (0,s+pads[i][1]) for i_ in range(x.ndim)])) + x.shrink(tuple([(0,s_) if i_ != i else (s_-1, s_) for i_,s_ in enumerate(x.shape)])).expand([pads[i][1] if i_ == i else s_ for i_,s_ in enumerate(x.shape)]).pad(tuple([(0,0) if i_ != i else (s+pads[i][0],0) for i_ in range(x.ndim)])) + x.pad(tuple([(0,0) if i_ != i else pads[i] for i_ in range(x.ndim)]))
|
||||
return x
|
||||
elif mode == "constant":
|
||||
return _padding(x, seq_pads, axes=seq_axes, constant_value=constant_value)
|
||||
|
||||
def AveragePool(X, kernel_shape, auto_pad="NOTSET", ceil_mode=0, count_include_pad=0, dilations=1, pads=None, strides=1):
|
||||
assert ceil_mode == 0 and dilations == 1
|
||||
if dilations != 1: raise NotImplementedError(f"dilations != 1 not supported, dilations:{dilations}")
|
||||
pixel_axes = tuple(range(len(X.shape)))[-2:]
|
||||
padding_included = _padding(X, pads, auto_pad, axes=pixel_axes).avg_pool2d(kernel_shape, stride=strides)
|
||||
if ceil_mode: auto_pad = "SAME_UPPER"
|
||||
padding_included = _padding(X, pads, auto_pad, axes=pixel_axes, strides=strides, kernel_shape=kernel_shape, dilations=dilations).avg_pool2d(kernel_shape, stride=strides)
|
||||
if count_include_pad:
|
||||
return padding_included
|
||||
else:
|
||||
div = _padding(Tensor.ones(*X.shape), pads, auto_pad, axes=pixel_axes).avg_pool2d(kernel_shape, stride=strides)
|
||||
div = _padding(Tensor.ones(*X.shape), pads, auto_pad, axes=pixel_axes, strides=strides, kernel_shape=kernel_shape, dilations=dilations).avg_pool2d(kernel_shape, stride=strides)
|
||||
return padding_included / div
|
||||
|
||||
def MaxPool(X, kernel_shape, auto_pad="NOTSET", ceil_mode=0, dilations=1, pads=None, storage_order=0, strides=1):
|
||||
assert ceil_mode == 0 and storage_order == 0, f"WARNING: MaxPool ceil_mode {ceil_mode} and storage_order {storage_order} not implemented"
|
||||
return _padding(X, pads, auto_pad, constant_value=-np.inf, axes=tuple(range(len(X.shape)))[-2:]).max_pool2d(kernel_shape, stride=strides, dilation=dilations)
|
||||
if ceil_mode: auto_pad = "SAME_UPPER"
|
||||
ret = _padding(X, pads, auto_pad, constant_value=-np.inf, axes=tuple(range(len(X.shape)))[-len(kernel_shape):], strides=strides, kernel_shape=kernel_shape, dilations=dilations)
|
||||
ret = ret.max_pool2d(kernel_shape, stride=strides, dilation=dilations)
|
||||
ret_len, X_len = ret.numel(), X.numel()
|
||||
indices = ((ret.flatten().unsqueeze(1).expand(ret_len, X_len) == X.flatten().reshape(1, X_len).expand(ret_len, X_len)) * Tensor.arange(X_len).reshape(1, X_len).expand(ret_len, X_len)).sum(1).reshape(ret.shape).cast(dtypes.int64)
|
||||
if storage_order: indices = indices.transpose(indices.ndim-2, indices.ndim-1)
|
||||
return ret, indices
|
||||
|
||||
def MaxUnpool(xT, xI, outshape=None, kernel_shape=None, pads=None, strides=None):
|
||||
out_sh = [(ks//2)*2 + st * inps for inps, st, ks in zip(xI.shape, strides, kernel_shape)]
|
||||
outlength = prod(out_sh)
|
||||
xI = xI.flatten().unsqueeze(1).expand(prod(xT.shape), outlength)
|
||||
arange = Tensor.arange(outlength).reshape(1, outlength).expand(xI.shape)
|
||||
xT = xT.flatten().unsqueeze(1).expand(prod(xT.shape), outlength)
|
||||
ret = ((xI == arange) * xT).sum(0).reshape([1, 1] + out_sh)
|
||||
if outshape is not None:
|
||||
outshape = safe_numpy(outshape).tolist()
|
||||
if outshape != ret.shape:
|
||||
diff = [outshape[2] - ret.shape[2], outshape[3] - ret.shape[3]]
|
||||
pad_args = [diff[0]//2, diff[1]//2, diff[0]-diff[0]//2, diff[1]-diff[1]//2]
|
||||
ret = ret.pad2d((pad_args[1], pad_args[3], pad_args[0], pad_args[2]))
|
||||
return ret
|
||||
|
||||
def Conv(X, W, B=None, auto_pad="NOTSET", dilations=1, group=1, kernel_shape=None, pads=None, strides=1):
|
||||
padding = [p for ps in zip(pads[:len(pads)//2][::-1], pads[len(pads)//2:][::-1]) for p in ps] if pads is not None else 0 # reorder padding
|
||||
if auto_pad != "NOTSET": padding = _auto_pad(X, auto_pad, strides, kernel_shape, dilations)
|
||||
else: padding = [p for ps in zip(pads[:len(pads)//2][::-1], pads[len(pads)//2:][::-1]) for p in ps] if pads is not None else 0 # reorder padding
|
||||
return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations, padding=padding)
|
||||
|
||||
def ConvTranspose(X, W, B=None, auto_pad="NOTSET", dilations=1, group=1, kernel_shape=None, pads=None, output_shape=None, output_padding=0, strides=1):
|
||||
return X.conv_transpose2d(W, B, stride=strides, groups=group, dilation=dilations, padding=(pads[1], pads[3], pads[0], pads[2]) if pads is not None else 0, output_padding=output_padding)
|
||||
if not kernel_shape: kernel_shape = W.shape
|
||||
if pads is None and auto_pad != "NOTSET": pads = _auto_pad(X, auto_pad, strides, kernel_shape, dilations)
|
||||
elif pads is None and auto_pad == "NOTSET": pads = [0,0] * (X.ndim - 2)
|
||||
strides_ = [1]*(W.ndim-1) + [strides] if isinstance(strides, int) else [1]*(W.ndim-len(strides)) + list(strides)
|
||||
dilations_ = [1]*(W.ndim-1) + [dilations] if isinstance(dilations, int) else [1]*(W.ndim-len(dilations)) + list(dilations)
|
||||
if output_shape and not output_padding:
|
||||
out_sh = [st*(xs-1) + (ks-1)*di+1 if n < 2 else st*(xs-1) + (ks-1)*di+1 - pads[n-2] - pads[n-1] for n, (st, xs, ks, di) in enumerate(zip(strides_, X.shape, kernel_shape, dilations_))]
|
||||
output_padding = [os - rs for os, rs in zip(output_shape, out_sh[-len(output_shape):])]
|
||||
return X.conv_transpose2d(W, B, stride=strides, groups=group, dilation=dilations, padding=pads if pads is not None else 0, output_padding=output_padding)
|
||||
|
||||
# Reimplemented here because you need legacy RNG for passing ONNX tests.
|
||||
def Dropout(data, ratio=0.5, training_mode=False, seed=None):
|
||||
if isinstance(ratio, Tensor) and not ratio.shape: ratio = safe_numpy(ratio) # ratio and tensor is passed in as Tensor with shape: ()
|
||||
if isinstance(training_mode, Tensor) and not training_mode.shape: training_mode = safe_numpy(training_mode)
|
||||
if not training_mode: return data, Tensor.ones(*data.shape, dtype=dtypes.bool) # if mask is requested as output it will contain all True's.
|
||||
rng = np.random.RandomState(seed)
|
||||
ratio = ratio.lazydata.realize().toCPU()[0] if isinstance(ratio, Tensor) else ratio
|
||||
@@ -113,11 +237,7 @@ def Dropout(data, ratio=0.5, training_mode=False, seed=None):
|
||||
|
||||
def Shape(data, end=None, start=0): return Tensor(list(data.shape)[start:end], dtype=dtypes.int64)
|
||||
def Size(data): return prod(data if isinstance(data, list) else data.shape)
|
||||
|
||||
# TODO: this doesn't match Tensor.flatten behavior
|
||||
def Flatten(input, axis=1):
|
||||
new_shape = (1, -1) if axis == 0 else (prod(input.shape[0:axis]), -1)
|
||||
return input.reshape(new_shape)
|
||||
def Flatten(input, axis=1): return input.reshape(prod((1,) + input.shape[0:axis]), -1)
|
||||
|
||||
# TODO: abstract out the broadcast logic in tensor
|
||||
def Expand(input, shape):
|
||||
@@ -146,7 +266,9 @@ def HardSwish(input): return input * HardSigmoid(input, 1/6, 0.5)
|
||||
def Celu(X, alpha=1.0): return X.relu() - (-alpha*(X/alpha).exp()+1).relu()
|
||||
def Selu(X, alpha=1.67326319217681884765625, gamma=1.05070102214813232421875): return gamma * (X.relu() - (-alpha*X.exp()+alpha).relu())
|
||||
def Softplus(X): return X.softplus()
|
||||
def PRelu(X:Tensor, slope:Tensor): return X.clip(0, float("inf")) + X.clip(float("-inf"), 0) * slope
|
||||
def PRelu(X:Tensor, slope:Tensor):
|
||||
slope = slope[0] if slope.shape[-1] != X.shape[-1] else slope # OnnxBackendPyTorchConvertedModelTest HAS WEIRD SLOPE WHERE IT'S [0.25, 0.25, 0.25] FOR ANY X.SHAPE
|
||||
return X.clip(0, float("inf")) + X.clip(float("-inf"), 0) * slope
|
||||
def LeakyRelu(X, alpha=0.01): return X.leakyrelu(alpha)
|
||||
def ThresholdedRelu(X, alpha=1.0): return (X-alpha).relu() + (X-alpha).relu().sign() * alpha
|
||||
def Softmax_1(input, axis=1): return input.softmax(axis)
|
||||
@@ -154,8 +276,8 @@ def Softmax_13(input, axis=-1): return input.softmax(axis)
|
||||
Softmax = {1: Softmax_1, 13: Softmax_13} # Softmax default axis changed
|
||||
def LogSoftmax(input, axis=-1): return input.log_softmax(axis)
|
||||
def Clip(input, min=None, max=None):
|
||||
if min is None: min = -3.4e38
|
||||
if max is None: max = 3.4e38
|
||||
if min is None: min = float("-inf")
|
||||
if max is None: max = float("inf")
|
||||
return input.clip(min, max)
|
||||
|
||||
def Sin(x): return x.sin()
|
||||
@@ -176,7 +298,8 @@ 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 functools.reduce(Tensor.__add__, data_0) / len(data_0)
|
||||
|
||||
def _axes(axes, noop_with_empty_axes): return [int(x) for x in safe_numpy(axes)] if axes is not None else ([] if noop_with_empty_axes else None)
|
||||
def _axes(axes, noop_with_empty_axes):
|
||||
return [int(x) for x in safe_numpy(axes)] if axes is not None and not (isinstance(axes, Tensor) and axes.shape == (0,)) else ([] if noop_with_empty_axes else None)
|
||||
|
||||
# ReduceProd would require a new llop
|
||||
def ReduceMax(data, axes=None, keepdims=1, noop_with_empty_axes=0): return data.max(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
|
||||
@@ -195,14 +318,9 @@ def GlobalMaxPool(X): return X.max(axis=tuple(range(2, len(X.shape))), keepdim=T
|
||||
def OptionalHasElement(x: Tensor=None): return Tensor(x is not None and x.numel() > 0, dtype=dtypes.bool)
|
||||
def OptionalGetElement(x: Tensor=None): return x if x is not None else Tensor([], dtype=dtypes.float32)
|
||||
|
||||
def Tile(input, repeats):
|
||||
repeats_ = [int(x) for x in safe_numpy(repeats)]
|
||||
new_shape = [x for i in range(len(input.shape)) for x in [1,input.shape[i]]]
|
||||
expand_shape = [x for r,s in zip(repeats_, input.shape) for x in [r,s]]
|
||||
final_shape = [r*s for r,s in zip(repeats_, input.shape)]
|
||||
return input.reshape(new_shape).expand(expand_shape).reshape(final_shape)
|
||||
def Tile(input, repeats): return input.repeat([int(x) for x in safe_numpy(repeats)])
|
||||
|
||||
def Range(start, limit, delta): return Tensor.arange(*[safe_numpy(x)[0].item() for x in (start, limit, delta)])
|
||||
def Range(start, limit, delta): return Tensor.arange(start=int(safe_numpy(start)), stop=int(safe_numpy(limit)), step=int(safe_numpy(delta))).cast(dtype=start.dtype) # DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated, and will error in future. Ensure you extract a single element from your array before performing this operation. (Deprecated NumPy 1.25.)
|
||||
def Where(condition:Tensor,X:Tensor,Y:Tensor): return condition.where(X, Y).cast(X.dtype)
|
||||
|
||||
def And(x:Tensor, y:Tensor): return Where((x==y), x, Tensor.zeros(*x.shape)).cast(dtypes.bool)
|
||||
@@ -210,6 +328,8 @@ def Or(x:Tensor, y:Tensor): return Where((x==y), x, Tensor.ones(*x.shape)).cast(
|
||||
def Xor(x:Tensor, y:Tensor): return Where((x==y), Tensor.zeros(*x.shape), Tensor.ones(*x.shape)).cast(dtypes.bool)
|
||||
def Not(x:Tensor): return Where((x==1), Tensor.zeros(*x.shape), Tensor.ones(*x.shape)).cast(dtypes.bool)
|
||||
|
||||
def Floor(x:Tensor): return x.floor()
|
||||
def Ceil(x:Tensor): return x.ceil()
|
||||
def Trilu(x: Tensor, k: Union[Tensor, int]=0, upper=1):
|
||||
k = int(k.numpy().item()) if k != 0 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)
|
||||
@@ -234,13 +354,14 @@ def MeanVarianceNormalization(input, axis=(0, 2, 3)):
|
||||
return (input - data_mean) / (std + 1e-9)
|
||||
|
||||
def NegativeLogLikelihoodLoss(input, target, weight=None, ignore_index=None, reduction="mean"):
|
||||
target = target.cast(dtypes.float32)
|
||||
N, C, i_shape = input.shape[0], input.shape[1], input.shape
|
||||
t_shape = target.shape
|
||||
if len(input.shape) != 3:
|
||||
input = input.reshape((N, C, -1))
|
||||
target = target.reshape((N, -1))
|
||||
if weight is not None:
|
||||
mask = target.unsqueeze(-1) == Tensor.arange(C,dtype=dtypes.int64).repeat((N, 1, 1))
|
||||
mask = target.unsqueeze(-1) == Tensor.arange(C).repeat((N, 1, 1))
|
||||
weight = (mask * weight).sum(axis=-1)
|
||||
if ignore_index is not None:
|
||||
cond = (target == ignore_index)
|
||||
@@ -251,16 +372,229 @@ def NegativeLogLikelihoodLoss(input, target, weight=None, ignore_index=None, red
|
||||
elif reduction == "sum": return loss.sum()
|
||||
return loss.reshape(t_shape) if len(i_shape) != 3 else loss
|
||||
|
||||
def SoftmaxCrossEntropyLoss(scores, labels, weights=None, ignore_index=None, reduction="mean"):
|
||||
N, C, *s_dimensions = scores.shape
|
||||
if ignore_index is not None: labels = (labels == ignore_index).where(C+1, labels)
|
||||
mask = labels.unsqueeze(1) == Tensor.arange(C).reshape(1, C, *[1]*len(s_dimensions))
|
||||
y = scores.log_softmax(axis=1)
|
||||
if weights is not None: weights = weights.__getitem__(tuple([labels, *[slice(None)]*(weights.ndim-1)]))
|
||||
loss = (mask * -y).sum(1) if weights is None else (mask * -y).sum(1) * weights
|
||||
if reduction == "mean": loss = loss.sum() / (loss == 0).where(0, 1).sum() if weights is None else loss.sum() / weights.sum()
|
||||
elif reduction == "sum": loss = loss.sum()
|
||||
return loss, y
|
||||
|
||||
def ArrayFeatureExtractor(input, indices): return input.__getitem__(tuple([slice(None) if i != (input.ndim-1) else indices for i in range(input.ndim)]))
|
||||
def Gather(input, indices, axis=0):
|
||||
if indices.numel() < 9: # NOTE lessor kernels for smaller indices but kernel number increases depending on size of indices
|
||||
input_sh = list(input.shape)
|
||||
ret_shape = input_sh[:axis] + list(indices.shape) + input_sh[axis+1:]
|
||||
if indices.ndim > 1: indices = indices.flatten()
|
||||
indices = [int(safe_numpy(indices))] if indices.shape == () else [input_sh[axis]+int(x) if x<0 else int(x) for x in safe_numpy(indices)]
|
||||
args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(input_sh)] for i in indices]
|
||||
return input.shrink(arg=tuple(args[0])).cat(*[input.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
|
||||
else: # NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
|
||||
return input.__getitem__(tuple([slice(None) if i != axis else indices for i in range(input.ndim)]))
|
||||
|
||||
def GatherElements(input, indices, axis):
|
||||
indices = indices.sign().contiguous().__neg__().contiguous().relu() * input.shape[axis] + indices
|
||||
return input.gather(indices, axis)
|
||||
|
||||
def _round(x:Tensor, n:float, equidistant_case = "round_down") -> Tensor:
|
||||
def _and(cond1, cond2): return ((cond1 + cond2) == 2).where(1, 0)
|
||||
assert n <= 1, f"n:{n} shouldn't be larger than 1"
|
||||
b = x.cast(dtypes.int32).contiguous().cast(x.dtype)
|
||||
b = (b >= 0).where(b+n, b-n)
|
||||
if equidistant_case == "round_down":
|
||||
return (x > b).where(b+1-n, b-n)
|
||||
elif equidistant_case == "round_up":
|
||||
return (x >= b).where(b+1-n, b-n)
|
||||
elif equidistant_case == "round_to_even":
|
||||
x_ceil_fraction = x.ceil()/2
|
||||
cond_ceil_even = x_ceil_fraction.ceil() == x_ceil_fraction
|
||||
x = (_and(x == b, cond_ceil_even)).where(x+1-n, x)
|
||||
x = (x > b).where(b+1-n, b-n)
|
||||
return x
|
||||
|
||||
def Round(X:Tensor): return _round(X, 0.5, "round_to_even")
|
||||
|
||||
def Resize(X:Tensor, roi=None, scales=None, sizes=None, antialias=0, axes=None, coordinate_transformation_mode='half_pixel', cubic_coeff_a=-0.75, exclude_outside=0, extrapolation_value=0.0, keep_aspect_ratio_policy='stretch', mode='nearest', nearest_mode='round_prefer_floor'):
|
||||
def _nearest_gather(X: Tensor, x_out, y_out): return X[:,:,y_out,:][:,:,:,x_out]
|
||||
def _nearest_mode(x_resized: Tensor, nearest_mode: str, x_len):
|
||||
if nearest_mode == "round_prefer_floor": ret = _round(x_resized, 0.5, "round_down")
|
||||
elif nearest_mode == "round_prefer_ceil": ret = _round(x_resized, 0.5, "round_up")
|
||||
elif nearest_mode == "floor": ret = x_resized.floor()
|
||||
elif nearest_mode == "ceil": ret = x_resized.ceil()
|
||||
return ret.clip(0, x_len-1)
|
||||
def _coordinate_transformation(x_out, y_out, output_shape, scales_lol, roi=None):
|
||||
if coordinate_transformation_mode == "half_pixel":
|
||||
x_out = (x_out + 0.5)/Tensor(scales_lol[-1]) - 0.5 # TODO Tensor() because try (((Tensor([0,1,2,3,4,5])+0.5)/3.5 - 0.5)) with LLVM or METAL, inaccuacy.
|
||||
y_out = (y_out + 0.5)/Tensor(scales_lol[-2]) - 0.5
|
||||
elif coordinate_transformation_mode == "align_corners":
|
||||
x_out = x_out * (X.shape[-1] - 1) / (output_shape[-1] - 1)
|
||||
y_out = y_out * (X.shape[-2] - 1) / (output_shape[-2] - 1)
|
||||
elif coordinate_transformation_mode == "asymmetric":
|
||||
x_out = x_out/scales_lol[-1]
|
||||
y_out = y_out/scales_lol[-2]
|
||||
elif coordinate_transformation_mode == "half_pixel_symmetric":
|
||||
x_out = X.shape[-1] / 2 * (1 - int(output_shape[-1]) / output_shape[-1]) + (x_out + 0.5) / scales_lol[-1] - 0.5
|
||||
y_out = X.shape[-2] / 2 * (1 - int(output_shape[-2]) / output_shape[-2]) + (y_out + 0.5) / scales_lol[-2] - 0.5
|
||||
elif coordinate_transformation_mode == "pytorch_half_pixel":
|
||||
x_out = (x_out + 0.5)/scales_lol[-1] - 0.5 if output_shape[-1] > 1 else Tensor([0])
|
||||
y_out = (y_out + 0.5)/scales_lol[-2] - 0.5 if output_shape[-2] > 1 else Tensor([0])
|
||||
elif coordinate_transformation_mode == "tf_crop_and_resize":
|
||||
x_out = roi[-1][0] * (X.shape[-1] - 1) + x_out * ((roi[-1][1] - roi[-1][0]) * (X.shape[-1] - 1) / (output_shape[-1] - 1)) if output_shape[-1] > 1 else Tensor([0.5 * (roi[-1][0] + roi[-1][1]) * (X.shape[-1] - 1)])
|
||||
y_out = roi[-2][0] * (X.shape[-2] - 1) + y_out * ((roi[-2][1] - roi[-2][0]) * (X.shape[-2] - 1) / (output_shape[-2] - 1)) if output_shape[-2] > 1 else Tensor([0.5 * (roi[-2][0] + roi[-2][1]) * (X.shape[-2] - 1)])
|
||||
return x_out.clip(0, X.shape[-1]-1), y_out.clip(0, X.shape[-2]-1)
|
||||
if roi is not None:
|
||||
roi = safe_numpy(roi)
|
||||
roi = [(st,ed) for st, ed in zip(roi[:len(roi)//2], roi[len(roi)//2:])]
|
||||
roi_ = [(1,1)] * 4
|
||||
if axes is not None:
|
||||
for a,r in zip(axes, roi):
|
||||
roi_[a] = r
|
||||
roi = roi_
|
||||
if scales is not None:
|
||||
scales = safe_numpy(scales).tolist()
|
||||
if axes is not None:
|
||||
scales_ = [1]*X.ndim
|
||||
for a,s in zip(axes, scales):
|
||||
scales_[a] = s
|
||||
scales = scales_
|
||||
elif sizes is not None:
|
||||
sizes = [int(i) for i in safe_numpy(sizes)]
|
||||
scales = []
|
||||
if axes is not None:
|
||||
sizes_ = [1]*X.ndim
|
||||
for a,s in zip(axes, sizes):
|
||||
sizes_[a] = s
|
||||
scales.append(s/X.shape[a])
|
||||
sizes = sizes_
|
||||
else: scales = [si/xs for xs, si in zip(X.shape, sizes)]
|
||||
if keep_aspect_ratio_policy == "not_larger":
|
||||
scale = min(scales)
|
||||
sizes = _round(Tensor(list(X.shape[-2:]))*scale, 0.5, "round_up")
|
||||
sizes = list(X.shape[:-2]) + [int(i) for i in safe_numpy(sizes)]
|
||||
elif keep_aspect_ratio_policy == "not_smaller":
|
||||
scale = max(scales)
|
||||
sizes = _round(Tensor(list(X.shape[-2:]))*scale, 0.5, "round_up")
|
||||
sizes = list(X.shape[:-2]) + [int(i) for i in safe_numpy(sizes)]
|
||||
output_shape = sizes if sizes else [math.floor(x*s) for x,s in zip(X.shape, scales)]
|
||||
output_shape_ = sizes if sizes else [x*s for x,s in zip(X.shape, scales)]
|
||||
scales_lol = [os/xs for xs, os in zip(X.shape, output_shape)]
|
||||
x_out = Tensor.arange(output_shape[-1])
|
||||
y_out = Tensor.arange(output_shape[-2])
|
||||
if mode == "nearest":
|
||||
x_out, y_out = _coordinate_transformation(x_out, y_out, output_shape, scales_lol, roi)
|
||||
x_out = _nearest_mode(x_out, nearest_mode, X.shape[-1])
|
||||
y_out = _nearest_mode(y_out, nearest_mode, X.shape[-1])
|
||||
return _nearest_gather(X, x_out, y_out)
|
||||
elif mode == "linear":
|
||||
x_out, y_out = _coordinate_transformation(x_out, y_out, output_shape_, scales, roi)
|
||||
ret = []
|
||||
for y in safe_numpy(y_out):
|
||||
for x in safe_numpy(x_out):
|
||||
x_floor, y_floor = int(x), int(y)
|
||||
y_shrink = (0, X.shape[2]) if X.shape[2] == 1 else (y_floor, y_floor+2) if y != y_floor else (y_floor, y_floor+1)
|
||||
x_shrink = (x_floor, x_floor+2) if x != x_floor else (x_floor, x_floor+1)
|
||||
shrink_args = ((0, X.shape[0]), (0, X.shape[1]), y_shrink, x_shrink)
|
||||
corners = safe_numpy(X.shrink(shrink_args))
|
||||
x1, x2, y1, y2 = x_floor, x_floor+1, y_floor, y_floor+1
|
||||
if x == x_floor and y == y_floor: # TODO https://en.wikipedia.org/wiki/Bilinear_interpolation#Weighted_mean maybe do weighted mean?
|
||||
ret.append(corners[0,0,0,0])
|
||||
elif x == x_floor:
|
||||
ret.append((corners[0,0,0,0] * (y2 - y) + corners[0,0,1,0] * (y - y1)) / (y2 - y1))
|
||||
elif y == y_floor:
|
||||
ret.append((corners[0,0,0,0] * (x2 - x) + corners[0,0,0,1] * (x - x1)) / (x2 - x1))
|
||||
else:
|
||||
ret.append((corners[0,0,0,0] * (x2 - x) * (y2 - y) + corners[0,0,0,1] * (x - x1) * (y2 - y) + corners[0,0,1,0] * (x2 - x) * (y - y1) + corners[0,0,1,1] * (x - x1) * (y - y1)) / ((x2 - x1) * (y2 - y1)))
|
||||
return Tensor(ret).reshape(output_shape)
|
||||
elif mode == "cubic":
|
||||
raise Exception("cubic interpolation is not implemented")
|
||||
|
||||
def CenterCropPad(input, shape, axes=None):
|
||||
if not axes: axes = list(range(input.ndim))
|
||||
shrink_arg = [(0,i) for i in input.shape]
|
||||
pad_arg = [(0,0) for _ in range(input.ndim)]
|
||||
shape = safe_numpy(shape).tolist()
|
||||
for s, x in zip(shape, axes):
|
||||
if s < input.shape[x]: shrink_arg[x] = (input.shape[x]//2 - s//2, input.shape[x]//2 + s//2) if s%2 == 0 else (input.shape[x]//2 - s//2 - 1, input.shape[x]//2 + s//2)
|
||||
elif s > input.shape[x]: pad_arg[x] = ((s - input.shape[x])//2, (s - input.shape[x])//2) if (s - input.shape[x])% 2 == 0 else ((s - input.shape[x])//2, (s - input.shape[x])//2 + 1)
|
||||
else: pass
|
||||
return input.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
|
||||
|
||||
def OneHot(indices, depth, values, axis=-1):
|
||||
depth = int(safe_numpy(depth).item())
|
||||
indices, rank = (indices.cast(dtypes.float32) < 0).where(indices+depth, indices), len(indices.shape)
|
||||
indices, rank = (indices < 0).where(indices+depth, indices), len(indices.shape)
|
||||
if axis < 0: axis += rank + 1
|
||||
ls, rs = indices.shape[0:axis], indices.shape[axis: rank]
|
||||
cond = indices[:,None] == Tensor.arange(depth).reshape((1,) * len(ls) + (depth,) + (1,) * len(rs))
|
||||
return cond.where(values[1], values[0]).cast(values.dtype)
|
||||
|
||||
def Floor(x:Tensor): return x.floor()
|
||||
def Ceil(x:Tensor): return x.ceil()
|
||||
def Erf(x):
|
||||
sign = x.sign()
|
||||
x = x.abs()
|
||||
t = 1.0 / (1.0 + 0.3275911 * x)
|
||||
term1 = 0.254829592 * t
|
||||
term2 = -0.284496736 * t ** 2
|
||||
term3 = 1.421413741 * t ** 3
|
||||
term4 = -1.453152027 * t ** 4
|
||||
term5 = 1.061405429 * t ** 5
|
||||
y = (term1 + term2 + term3 + term4 + term5)
|
||||
return sign * (1.0 - y * Tensor.exp(-x * x))
|
||||
|
||||
def Compress(inp, condition, axis=None):
|
||||
if axis == None:
|
||||
inp = inp.flatten()
|
||||
axis = 0
|
||||
|
||||
axis = axis + inp.ndim if axis < 0 else axis
|
||||
|
||||
con_np = condition.numpy()
|
||||
con = Tensor(np.arange(condition.shape[0])[con_np]) # no boolean indexing in Tensor
|
||||
return inp.__getitem__(tuple([slice(None) if i != axis else con for i in range(inp.ndim)]))
|
||||
|
||||
def Acos(x):
|
||||
negate = (x < 0)
|
||||
x = x.abs()
|
||||
ret = ((((-0.0187293 * x) + 0.0742610)*x - 0.2121144) * x + 1.5707288) * Tensor.sqrt(1.0 - x)
|
||||
ret = ret - 2 * negate * ret
|
||||
return negate * 3.14159265358979 + ret
|
||||
|
||||
|
||||
def Atan(y):
|
||||
x = Tensor.ones(y.shape)
|
||||
t3 = x
|
||||
t1 = y.abs()
|
||||
t0 = (t3 > t1).where(t3, t1)
|
||||
t1 = (t3 < t1).where(t3, t1)
|
||||
t3 = t1 / t0
|
||||
t4 = t3 * t3
|
||||
t0 = ((((-0.013480470 * t4 + 0.057477314) * t4 - 0.121239071) * t4 + 0.195635925) * t4 - 0.332994597) * t4 + 0.999995630
|
||||
t3 = t0 * t3
|
||||
t3 = (y.abs() > x.abs()).where(1.570796327 - t3, t3)
|
||||
return (y < 0).where(-t3, t3)
|
||||
|
||||
def Asin(x): return Atan(x / Tensor.sqrt(1 - x * x))
|
||||
|
||||
def Asinh(x): return Tensor.log(x + Tensor.sqrt(x * x + 1))
|
||||
def Acosh(x): return Tensor.log(x + Tensor.sqrt(x * x - 1))
|
||||
def Atanh(x): return 0.5 * Tensor.log((1 + x)/(1 - x))
|
||||
|
||||
# Needs work
|
||||
def IsInf(x,detect_negative=1,detect_positive=1):
|
||||
ret = (x == float("inf"))*detect_positive + (x == float("-inf"))*detect_negative + Tensor.zeros(*x.shape)
|
||||
return ret.cast(dtypes.bool)
|
||||
|
||||
# Needs work
|
||||
def DequantizeLinear(x, x_scale, x_zero_point=0, axis=1):
|
||||
axis = axis + x.ndim if axis < 0 else axis
|
||||
x_sc = x_scale.reshape(*[1]*axis, *x_scale.shape, *[1]*(x.ndim - axis - x_scale.ndim))
|
||||
x_zer = x_zero_point.reshape(*[1]*(axis), *x_scale.shape, *[1]*(x.ndim - axis - x_scale.ndim)) if isinstance(x_zero_point, Tensor) else x_zero_point
|
||||
return ((x - x_zer) * x_sc)
|
||||
|
||||
# Needs work
|
||||
def IsNaN(x):
|
||||
return (x < float("-inf")).cast(dtypes.bool)
|
||||
|
||||
def EmbedLayerNormalization(input_ids, segment_ids:Optional[Tensor]=None, word_embedding:Tensor=None, position_embedding:Tensor=None, segment_embedding:Optional[Tensor]=None, gamma=None, beta=None, mask:Optional[Tensor]=None, position_ids:Optional[Tensor]=None, epsilon=None, mask_index_type=None):
|
||||
# https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.EmbedLayerNormalization
|
||||
@@ -328,3 +662,26 @@ def SkipLayerNormalization(input:Tensor, skip:Tensor, gamma, beta:Optional[Tenso
|
||||
def FastGelu(x:Tensor, bias:Optional[Tensor]=None):
|
||||
x = x + bias
|
||||
return 0.5 * x * (1 + (x * 0.797885 + 0.035677 * x ** 3).tanh())
|
||||
|
||||
def ArgMax(x, axis=0, keepdims=1, select_last_index=0):
|
||||
axis = axis + x.ndim if axis < 0 else axis
|
||||
m = x == (x.max(axis=axis, keepdim=keepdims) if keepdims else x.max(axis=axis, keepdim=keepdims).unsqueeze(axis))
|
||||
c = Tensor.arange(x.shape[axis]).reshape(*[1]*(axis), x.shape[axis], *[1]*(x.ndim - axis-1)) * m
|
||||
return c.max(axis=axis,keepdim=keepdims).cast(dtypes.int64)
|
||||
|
||||
def ArgMin(x, axis=0, keepdims=1, select_last_index=0): return ArgMax(-x, axis=axis, keepdims=keepdims, select_last_index=select_last_index)
|
||||
|
||||
def Upsample(X, scales, mode):
|
||||
return Resize(X=X, scales=scales, mode=mode)
|
||||
|
||||
type_map = {TensorProto.DOUBLE: dtypes.double, TensorProto.FLOAT: dtypes.float32}
|
||||
def EyeLike(x, dtype=None, k=0):
|
||||
if dtype is None: dtype = x.dtype
|
||||
else: dtype = type_map[dtype]
|
||||
shape = x.shape
|
||||
dim = min(x.shape)
|
||||
if shape[0] == shape[1]: return Tensor.eye(dim=dim, dtype=dtype)
|
||||
else:
|
||||
diff = (shape[0]-dim, shape[1]-dim)
|
||||
padarg = tuple([(d, d) if d == 0 else (k, d-k) for d in diff])
|
||||
return Tensor.eye(dim=dim, dtype=dtype).pad(padarg)
|
||||
|
||||
3
test/external/external_model_benchmark.py
vendored
3
test/external/external_model_benchmark.py
vendored
@@ -14,8 +14,7 @@ from tinygrad.ops import Device
|
||||
|
||||
MODELS = {
|
||||
"resnet50": "https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-caffe2-v1-9.onnx",
|
||||
# numerical issue with v0.9.4
|
||||
# "openpilot": "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx",
|
||||
"openpilot": "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx",
|
||||
"efficientnet": "https://github.com/onnx/models/raw/main/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx",
|
||||
"shufflenet": "https://github.com/onnx/models/raw/main/vision/classification/shufflenet/model/shufflenet-9.onnx",
|
||||
"commavq": "https://huggingface.co/commaai/commavq-gpt2m/resolve/main/gpt2m.onnx",
|
||||
|
||||
127
test/external/external_test_onnx_backend.py
vendored
127
test/external/external_test_onnx_backend.py
vendored
@@ -20,7 +20,7 @@ class TinygradModel(BackendRep):
|
||||
def run(self, inputs: Any, **kwargs: Any) -> Tuple[Any, ...]:
|
||||
real_inputs = {k:v for k,v in zip(self.input_names, inputs)}
|
||||
ret = self.fxn(real_inputs, debug=True)
|
||||
return tuple(x.numpy() if isinstance(x, Tensor) else np.array(x) for x in ret.values())
|
||||
return tuple(x.numpy() if isinstance(x, Tensor) else [i.numpy() for i in x] if isinstance(x, list) else np.array(x) for x in ret.values())
|
||||
|
||||
class TinygradBackend(Backend):
|
||||
@classmethod
|
||||
@@ -38,20 +38,13 @@ class TinygradBackend(Backend):
|
||||
|
||||
backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)
|
||||
|
||||
# add support for SoftmaxCrossEntropyLoss and NegativeLogLikelihoodLoss
|
||||
backend_test.exclude('test_sce_*')
|
||||
|
||||
# no support for reduce with multiply (needs llop)
|
||||
backend_test.exclude('test_reduce_prod_*')
|
||||
|
||||
# no optimizers (add them?)
|
||||
backend_test.exclude('test_adagrad_*')
|
||||
backend_test.exclude('test_adam_*')
|
||||
backend_test.exclude('test_nesterov_momentum_*')
|
||||
backend_test.exclude('test_momentum_*')
|
||||
|
||||
# disable some creation ops
|
||||
backend_test.exclude('test_eyelike_*')
|
||||
# TODO figure out why it's returning wrong values, geohotstan's uneducated guess is it's due to imprecision from float64 (double) -> float32
|
||||
# see Type Constraints: https://onnx.ai/onnx/operators/onnx_aionnxpreviewtraining_Adam.html#type-constraints
|
||||
backend_test.exclude('test_adam_multiple_cpu')
|
||||
backend_test.exclude('test_nesterov_momentum_cpu')
|
||||
|
||||
# we only support float32
|
||||
backend_test.exclude('uint8')
|
||||
@@ -69,40 +62,25 @@ backend_test.exclude('test_castlike_*')
|
||||
backend_test.exclude('test_convinteger_*')
|
||||
backend_test.exclude('test_matmulinteger_*')
|
||||
|
||||
# we don't support rounding
|
||||
backend_test.exclude('test_round_*')
|
||||
backend_test.exclude('test_reduce_log_sum_exp*') # dependent on actual float64 implementation for backends
|
||||
backend_test.exclude('test_operator_add*') # dependent on float64 math. Without it values default to 0 or inf
|
||||
|
||||
# we don't support indexes
|
||||
backend_test.exclude('test_argmax_*')
|
||||
backend_test.exclude('test_argmin_*')
|
||||
# backend_test.exclude('test_argmax_*') # Needs more work #select_last_index
|
||||
# backend_test.exclude('test_argmin_*') # Needs more work #select_last_index
|
||||
backend_test.exclude('test_nonzero_*')
|
||||
|
||||
# no support for nan or inf
|
||||
backend_test.exclude('test_isinf_*')
|
||||
backend_test.exclude('test_isnan_*')
|
||||
|
||||
# no support for mod
|
||||
backend_test.exclude('test_mod_*')
|
||||
|
||||
# no trig ops
|
||||
backend_test.exclude('test_acos_*')
|
||||
backend_test.exclude('test_acosh_*')
|
||||
backend_test.exclude('test_asin_*')
|
||||
backend_test.exclude('test_asinh_*')
|
||||
backend_test.exclude('test_atan_*')
|
||||
backend_test.exclude('test_atanh_*')
|
||||
|
||||
# no boolean ops (2d, 3d, 4d)
|
||||
backend_test.exclude('test_bitshift_*')
|
||||
|
||||
# no scatter gather
|
||||
backend_test.exclude('test_gather_*')
|
||||
# no scatternd gathernd
|
||||
backend_test.exclude('test_gathernd_*')
|
||||
backend_test.exclude('test_scatter_*')
|
||||
backend_test.exclude('test_scatternd_*')
|
||||
|
||||
# no quantize
|
||||
backend_test.exclude('test_dequantizelinear_*')
|
||||
backend_test.exclude('test_dynamicquantizelinear_*')
|
||||
backend_test.exclude('test_qlinearmatmul_*')
|
||||
backend_test.exclude('test_qlinearconv_*')
|
||||
@@ -117,12 +95,14 @@ backend_test.exclude('test_simple_rnn_*')
|
||||
# no control flow
|
||||
backend_test.exclude('test_if_*')
|
||||
backend_test.exclude('test_loop*')
|
||||
backend_test.exclude('test_range_float_type_positive_delta_expanded_cpu') # requires loop
|
||||
|
||||
# unsupported (strange) ops
|
||||
backend_test.exclude('test_bitwise_*')
|
||||
backend_test.exclude('test_blackmanwindow_*')
|
||||
backend_test.exclude('test_bernoulli_*')
|
||||
backend_test.exclude('test_cumsum_*')
|
||||
backend_test.exclude('test_det_*')
|
||||
|
||||
backend_test.exclude('test_tril_zero_cpu') # TODO: zero array support
|
||||
backend_test.exclude('test_triu_zero_cpu') # TODO: zero array support
|
||||
@@ -132,11 +112,8 @@ backend_test.exclude('test_hammingwindow_*')
|
||||
backend_test.exclude('test_hannwindow_*')
|
||||
backend_test.exclude('test_hardmax_*')
|
||||
backend_test.exclude('test_gridsample_*')
|
||||
backend_test.exclude('test_compress_*')
|
||||
backend_test.exclude('test_det_*')
|
||||
backend_test.exclude('test_dft_*')
|
||||
backend_test.exclude('test_einsum_*')
|
||||
backend_test.exclude('test_erf_*')
|
||||
backend_test.exclude('test_strnorm_*')
|
||||
backend_test.exclude('test_unique_*')
|
||||
backend_test.exclude('test_sequence_*')
|
||||
@@ -149,15 +126,91 @@ backend_test.exclude('test_stft_*')
|
||||
backend_test.exclude('test_melweightmatrix_*')
|
||||
|
||||
# more strange ops
|
||||
backend_test.exclude('test_center_crop_pad_crop_*')
|
||||
backend_test.exclude('test_basic_deform_conv_*')
|
||||
backend_test.exclude('test_deform_conv_*')
|
||||
backend_test.exclude('test_lppool_*')
|
||||
backend_test.exclude('test_depthtospace_*')
|
||||
backend_test.exclude('test_spacetodepth_*')
|
||||
backend_test.exclude('test_scan*')
|
||||
backend_test.exclude('test_ai_onnx_ml_array_feature_extractor_*')
|
||||
backend_test.exclude('test_split_to_sequence_*')
|
||||
backend_test.exclude('test_resize_downsample_scales_cubic_*') # unsure how to implement cubic
|
||||
backend_test.exclude('test_resize_downsample_sizes_cubic_*') # unsure how to implement cubic
|
||||
backend_test.exclude('test_resize_upsample_scales_cubic_*') # unsure how to implement cubic
|
||||
backend_test.exclude('test_resize_upsample_sizes_cubic_*') # unsure how to implement cubic
|
||||
|
||||
# rest of the failing tests
|
||||
backend_test.exclude('test_averagepool_2d_dilations_cpu') # dilations != 1 not supported for avgpool
|
||||
backend_test.exclude('test_convtranspose_autopad_same_cpu') # TODO geohotstan has no idea how this is done, autopad requires output_shape but output_shape requires pads from autopad
|
||||
backend_test.exclude('test_optional_has_element_empty_optional_input_cpu') # Attempts to create Tensor from None
|
||||
backend_test.exclude('test_range_int32_type_negative_delta_expanded_cpu') # AttributeProto.GRAPH not implemented
|
||||
backend_test.exclude('test_reshape_allowzero_reordered_cpu') # reshaping to 0 shape
|
||||
backend_test.exclude('test_resize_downsample_scales_linear_antialias_cpu') # antialias not implemented
|
||||
backend_test.exclude('test_resize_downsample_sizes_linear_antialias_cpu') # antialias not implemented
|
||||
backend_test.exclude('test_resize_tf_crop_and_resize_cpu') # unsure about fill value after clip
|
||||
backend_test.exclude('test_operator_addconstant_cpu') # bad data type
|
||||
|
||||
# 1556
|
||||
backend_test.exclude('test_isinf_cpu')
|
||||
backend_test.exclude('test_isinf_negative_cpu')
|
||||
backend_test.exclude('test_isinf_positive_cpu')
|
||||
backend_test.exclude('test_isnan_cpu')
|
||||
|
||||
if getenv("CPU") or getenv("ARM64"):
|
||||
# not too sure
|
||||
backend_test.exclude('test_dequantizelinear_axis_cpu')
|
||||
backend_test.exclude('test_dequantizelinear_cpu')
|
||||
|
||||
if getenv("TORCH"): # 1562
|
||||
backend_test.exclude('test_and2d_cpu')
|
||||
backend_test.exclude('test_and3d_cpu')
|
||||
backend_test.exclude('test_and4d_cpu')
|
||||
backend_test.exclude('test_and_bcast3v1d_cpu')
|
||||
backend_test.exclude('test_and_bcast3v2d_cpu')
|
||||
backend_test.exclude('test_and_bcast4v2d_cpu')
|
||||
backend_test.exclude('test_and_bcast4v3d_cpu')
|
||||
backend_test.exclude('test_and_bcast4v4d_cpu')
|
||||
backend_test.exclude('test_dequantizelinear_axis_cpu')
|
||||
backend_test.exclude('test_dequantizelinear_cpu')
|
||||
backend_test.exclude('test_greater_equal_bcast_expanded_cpu')
|
||||
backend_test.exclude('test_greater_equal_expanded_cpu')
|
||||
backend_test.exclude('test_isinf_cpu')
|
||||
backend_test.exclude('test_isinf_negative_cpu')
|
||||
backend_test.exclude('test_isinf_positive_cpu')
|
||||
backend_test.exclude('test_isnan_cpu')
|
||||
backend_test.exclude('test_less_equal_bcast_expanded_cpu')
|
||||
backend_test.exclude('test_less_equal_expanded_cpu')
|
||||
backend_test.exclude('test_or2d_cpu')
|
||||
backend_test.exclude('test_or3d_cpu')
|
||||
backend_test.exclude('test_or4d_cpu')
|
||||
backend_test.exclude('test_or_bcast3v1d_cpu')
|
||||
backend_test.exclude('test_or_bcast3v2d_cpu')
|
||||
backend_test.exclude('test_or_bcast4v2d_cpu')
|
||||
backend_test.exclude('test_or_bcast4v3d_cpu')
|
||||
backend_test.exclude('test_or_bcast4v4d_cpu')
|
||||
backend_test.exclude('test_xor2d_cpu')
|
||||
backend_test.exclude('test_xor3d_cpu')
|
||||
backend_test.exclude('test_xor4d_cpu')
|
||||
backend_test.exclude('test_xor_bcast3v1d_cpu')
|
||||
backend_test.exclude('test_xor_bcast3v2d_cpu')
|
||||
backend_test.exclude('test_xor_bcast4v2d_cpu')
|
||||
backend_test.exclude('test_xor_bcast4v3d_cpu')
|
||||
backend_test.exclude('test_xor_bcast4v4d_cpu')
|
||||
|
||||
if getenv('LLVM') or getenv('GPU') or getenv('CLANG') or getenv('METAL') or getenv('MPS'):
|
||||
# compiled backends cannot reshape to 0 or from 0
|
||||
backend_test.exclude('test_slice_start_out_of_bounds_cpu')
|
||||
backend_test.exclude('test_constantofshape_int_shape_zero_cpu')
|
||||
|
||||
if getenv('GPU') or getenv('METAL') or getenv('MPS'):
|
||||
backend_test.exclude('test_mish_cpu') # weird inaccuracy
|
||||
backend_test.exclude('test_mish_expanded_cpu') # weird inaccuracy
|
||||
backend_test.exclude('test_eyelike_with_dtype_cpu') # I'm not sure about this...
|
||||
|
||||
if getenv('METAL') or getenv('MPS'):
|
||||
# (((Tensor([0,1,2,3,4,5])+0.5)/3.5 - 0.5)) Try this with METAL and LLVM, weird weird inaccuracy
|
||||
backend_test.exclude('test_resize_upsample_sizes_nearest_axes_2_3_cpu')
|
||||
backend_test.exclude('test_resize_upsample_sizes_nearest_axes_3_2_cpu')
|
||||
backend_test.exclude('test_resize_upsample_sizes_nearest_cpu')
|
||||
|
||||
# disable model tests for now since they are slow
|
||||
if not getenv("MODELTESTS"):
|
||||
|
||||
@@ -16,7 +16,6 @@ base_fxn_for_op: Dict[Op, Callable] = {
|
||||
MovementOps.RESHAPE: lambda x, arg: x.reshape(arg), MovementOps.SHRINK: lambda x, arg: x[tuple(slice(p[0], p[1], None) for p in arg)],
|
||||
}
|
||||
|
||||
def promote_types(x, y): return ret if (ret := np.promote_types(x.dtype, y.dtype)) != np.float64 else np.float32
|
||||
def match_types(x, y):
|
||||
up = x.dtype if dtypes.from_np(x.dtype).priority > dtypes.from_np(y.dtype).priority else y.dtype
|
||||
return x.astype(up, copy=False), y.astype(up, copy=False)
|
||||
@@ -34,7 +33,7 @@ def einsum_mulacc(einsum, get_strides, expand):
|
||||
numpy_fxn_for_op: Dict[Op, Callable] = {**base_fxn_for_op, **{
|
||||
UnaryOps.NOOP: lambda x: np.require(x, requirements='C'), UnaryOps.EXP2: np.exp2, UnaryOps.LOG2: np.log2, UnaryOps.SIN: np.sin,
|
||||
UnaryOps.CAST: lambda x,y: x.view(y[0].np) if y[1] else x.astype(y[0].np, copy=False),
|
||||
BinaryOps.MAX: np.maximum, BinaryOps.CMPLT: lambda x,y: (x<y).astype(promote_types(x,y)), BinaryOps.ADD: lambda x, y: np.add(*match_types(x, y)),
|
||||
BinaryOps.MAX: np.maximum, BinaryOps.CMPLT: lambda x,y: (x<y).astype(np.promote_types(x.dtype,y.dtype)), BinaryOps.ADD: lambda x, y: np.add(*match_types(x, y)),
|
||||
BinaryOps.SUB: lambda x, y: np.subtract(*match_types(x, y)), BinaryOps.MUL: lambda x, y: np.multiply(*match_types(x, y)),
|
||||
BinaryOps.DIV: lambda x, y: np.divide(*match_types(x, y)), UnaryOps.SQRT: np.sqrt,
|
||||
MovementOps.PERMUTE: lambda x, order: x.transpose(order), MovementOps.PAD: np.pad, MovementOps.EXPAND: np.broadcast_to,
|
||||
|
||||
Reference in New Issue
Block a user