forked from tinygrad/tinygrad
onnx_ops formatting cleanup (#2904)
also removed a case in safe_numpy that always convert 0-dim array to 1-dim
This commit is contained in:
chenyu
+1
-2
@@ -21,8 +21,7 @@ def safe_numpy(t) -> np.ndarray:
|
||||
if t not in numpy_cache:
|
||||
if DEBUG >= 3: print("numpy cache miss", t)
|
||||
tmp = t.numpy()
|
||||
numpy_cache[t] = tmp if len(tmp.shape) else tmp.reshape(1)
|
||||
assert len(numpy_cache[t].shape) > 0
|
||||
numpy_cache[t] = tmp
|
||||
return numpy_cache[t]
|
||||
|
||||
onnx_ops = importlib.import_module('extra.onnx_ops')
|
||||
|
||||
+107
-114
@@ -13,11 +13,11 @@ tensor_methods = {"Neg", "Reciprocal", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mis
|
||||
|
||||
# **************** Free Ops ****************
|
||||
|
||||
def Identity(input: Tensor): return input
|
||||
def Add(input: Tensor, other: Tensor, broadcast=None): return input + other if input.dtype == dtypes.float or isinstance(input.dtype, ImageDType) else (input + other).cast(input.dtype)
|
||||
def Sub(input: Union[Tensor, Any], other: Tensor): return input - other # some test has input as int
|
||||
def Div(input: Tensor, other: Tensor): return input / other if input.dtype == dtypes.float or isinstance(input.dtype, ImageDType) else input.div(other).floor() # TODO: this has dtype issues
|
||||
def Pow(input: Tensor, other: Tensor): return input.float() ** other.float()
|
||||
def Identity(x: Tensor): return x
|
||||
def Add(x: Tensor, other: Tensor, broadcast=None): return x + other if x.dtype == dtypes.float or isinstance(x.dtype, ImageDType) else (x + other).cast(x.dtype)
|
||||
def Sub(x: Union[Tensor, Any], other: Tensor): return x - other # some test has input as int
|
||||
def Div(x: Tensor, other: Tensor): return x / other if x.dtype == dtypes.float or isinstance(x.dtype, ImageDType) else x.div(other).floor() # TODO: this has dtype issues
|
||||
def Pow(x: Tensor, other: Tensor): return x.float() ** other.float()
|
||||
def Less(x:Tensor,y:Tensor): return (x<y).cast(dtypes.bool)
|
||||
def LessOrEqual(x:Tensor,y:Tensor): return (x<=y).cast(dtypes.bool)
|
||||
def Greater(x:Tensor,y:Tensor): return (x>y).cast(dtypes.bool)
|
||||
@@ -27,19 +27,19 @@ def Max(*data_0): return functools.reduce(Tensor.maximum, data_0)
|
||||
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 Cast(input: Tensor, to): return input.cast(dtypes.from_np(tensor_dtype_to_np_dtype(to)))
|
||||
def Cast(x: Tensor, to): return x.cast(dtypes.from_np(tensor_dtype_to_np_dtype(to)))
|
||||
|
||||
# **************** Simple Ops ****************
|
||||
|
||||
def Constant(value: Tensor=None, value_float=None, value_floats=None, value_int=None, value_ints=None, value_string=None, value_strings=None):
|
||||
if value: return value
|
||||
elif value_float: return Tensor(value_float, dtype=dtypes.float32, requires_grad=False)
|
||||
elif value_floats: return Tensor(list(value_floats), dtype=dtypes.float32, requires_grad=False)
|
||||
elif value_int: return Tensor(value_int, dtype=dtypes.int64, requires_grad=False)
|
||||
elif value_ints: return Tensor(list(value_ints), dtype=dtypes.int64, requires_grad=False)
|
||||
elif value_string or value_strings: raise NotImplementedError(f'value_string or value_strings not implemented for Constant op')
|
||||
if value_float: return Tensor(value_float, dtype=dtypes.float32, requires_grad=False)
|
||||
if value_floats: return Tensor(list(value_floats), dtype=dtypes.float32, requires_grad=False)
|
||||
if value_int: return Tensor(value_int, dtype=dtypes.int64, requires_grad=False)
|
||||
if value_ints: return Tensor(list(value_ints), dtype=dtypes.int64, requires_grad=False)
|
||||
if value_string or value_strings: raise NotImplementedError('value_string or value_strings not implemented for Constant op')
|
||||
|
||||
def HardSigmoid(input: Tensor, alpha=0.2, beta=0.5): return (alpha*input + beta).clip(0, 1)
|
||||
def HardSigmoid(x: Tensor, alpha=0.2, beta=0.5): return (alpha*x + beta).clip(0, 1)
|
||||
def Gelu(x:Tensor, approximate=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + Erf(x/math.sqrt(2)))
|
||||
def Selu(X: Tensor, alpha=1.67326319217681884765625, gamma=1.05070102214813232421875): return gamma * (X.relu() - (-alpha*X.exp()+alpha).relu())
|
||||
def PRelu(X:Tensor, slope:Tensor):
|
||||
@@ -47,11 +47,11 @@ def PRelu(X:Tensor, slope:Tensor):
|
||||
return (X > 0).where(X, X * slope)
|
||||
def LeakyRelu(X: Tensor, alpha=0.01): return X.leakyrelu(alpha)
|
||||
def ThresholdedRelu(X: Tensor, alpha=1.0): return (X > alpha).where(X, 0)
|
||||
def Softmax_1(input: Tensor, axis=1): return input.softmax(axis)
|
||||
def Softmax_13(input: Tensor, axis=-1): return input.softmax(axis)
|
||||
def Softmax_1(x: Tensor, axis=1): return x.softmax(axis)
|
||||
def Softmax_13(x: Tensor, axis=-1): return x.softmax(axis)
|
||||
Softmax = {1: Softmax_1, 13: Softmax_13} # Softmax default axis changed
|
||||
def LogSoftmax(input: Tensor, axis=-1): return input.log_softmax(axis)
|
||||
def Clip(input: Tensor, min=None, max=None): return input.clip(float('-inf') if min is None else min, float('inf') if max is None else max)
|
||||
def LogSoftmax(x: Tensor, axis=-1): return x.log_softmax(axis)
|
||||
def Clip(x: Tensor, min=None, max=None): return x.clip(float('-inf') if min is None else min, float('inf') if max is None else max)
|
||||
|
||||
# NOTE ReduceProd would require a new llop
|
||||
def _axes(axes, noop_with_empty_axes):
|
||||
@@ -72,14 +72,14 @@ def GlobalMaxPool(X: Tensor): return X.max(axis=tuple(range(2, len(X.shape))), k
|
||||
def OptionalHasElement(x: Optional[Tensor]=None): return Tensor(x is not None and x.numel() > 0)
|
||||
def OptionalGetElement(x: Optional[Tensor]=None): return x if x is not None else Tensor([])
|
||||
|
||||
def Tile(input: Tensor, repeats): return input.repeat([int(x) for x in safe_numpy(repeats)])
|
||||
def Tile(x: Tensor, repeats): return x.repeat([int(x) for x in safe_numpy(repeats)])
|
||||
def Range(start: Tensor, limit, delta): return Tensor.arange(start=safe_numpy(start).item(), stop=safe_numpy(limit).item(), step=safe_numpy(delta).item())
|
||||
def Shape(data: Tensor, end=None, start=0): return Tensor(list(data.shape)[start:end], dtype=dtypes.int64)
|
||||
def Size(data: Tensor): return prod(data if isinstance(data, list) else data.shape)
|
||||
def Flatten(input: Tensor, axis=1): return input.reshape(prod((1,) + input.shape[0:axis]), -1)
|
||||
def Flatten(x: Tensor, axis=1): return x.reshape(prod((1,) + x.shape[0:axis]), -1)
|
||||
def Reshape(data: Tensor, shape: Tensor, allowzero=0):
|
||||
return data.reshape([int(x) if x != 0 else (0 if allowzero else data.shape[i]) for i,x in enumerate(safe_numpy(shape))])
|
||||
def Shrink(input: Tensor, bias=0.0, lambd=0.5): return (input < -lambd)*(input+bias) + (input > lambd)*(input-bias)
|
||||
def Shrink(x: Tensor, bias=0.0, lambd=0.5): return (x < -lambd)*(x+bias) + (x > lambd)*(x-bias)
|
||||
def And(x:Tensor, y:Tensor): return (x==y).where(x, 0).cast(dtypes.bool)
|
||||
def Or(x:Tensor, y:Tensor): return (x==y).where(x, 1).cast(dtypes.bool)
|
||||
def Xor(x:Tensor, y:Tensor): return (x==y).where(0, 1).cast(dtypes.bool)
|
||||
@@ -122,7 +122,7 @@ def Unsqueeze(data: Tensor, axes):
|
||||
new_shape[i] = next(ptr)
|
||||
return data.reshape(new_shape)
|
||||
|
||||
def Binarizer(input, threshold=0.0): return (input > threshold).cast(dtypes.float32)
|
||||
def Binarizer(x, threshold=0.0): return (x > threshold).cast(dtypes.float32)
|
||||
|
||||
def ArgMax(x: Tensor, axis=0, keepdims=1, select_last_index=0):
|
||||
axis = axis + x.ndim if axis < 0 else axis
|
||||
@@ -131,27 +131,27 @@ def ArgMax(x: Tensor, axis=0, keepdims=1, select_last_index=0):
|
||||
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 Concat(*inputs: List[Tensor], axis): return inputs[0].cat(*inputs[1:], dim=axis)
|
||||
def Transpose(input: Tensor, perm=None): return input.permute(order=list(range(len(input.shape))[::-1]) if perm is None else perm)
|
||||
def Concat(*xs: List[Tensor], axis): return xs[0].cat(*xs[1:], dim=axis)
|
||||
def Transpose(x: Tensor, perm=None): return x.permute(order=list(range(len(x.shape))[::-1]) if perm is None else perm)
|
||||
|
||||
# NOTE: since we only have one type, this is valid!
|
||||
# TODO: fix this with dtypes
|
||||
def CastLike(input, target_type):
|
||||
def CastLike(x, target_type):
|
||||
assert isinstance(target_type, Tensor), "can only CastLike Tensor"
|
||||
return input
|
||||
return x
|
||||
|
||||
def ConstantOfShape(input, value:Tensor=None):
|
||||
def ConstantOfShape(x, value:Tensor=None):
|
||||
if value is None: value=Tensor([0.0])
|
||||
shape = [int(x) for x in safe_numpy(input)]
|
||||
shape = [int(x) for x in safe_numpy(x)]
|
||||
return Tensor.ones(*shape, dtype=value.dtype) * (value if shape[0]!=0 else 1)
|
||||
|
||||
# TODO: abstract out the broadcast logic in tensor
|
||||
def Expand(input: Tensor, shape):
|
||||
x_shape, y_shape = input.shape, [int(x) for x in safe_numpy(shape)]
|
||||
def Expand(x: Tensor, shape):
|
||||
x_shape, y_shape = x.shape, [int(x) for x in safe_numpy(shape)]
|
||||
# copied from _broadcasted
|
||||
x_shape, y_shape = [([1]*(max(len(x_shape), len(y_shape))-len(t_shape)) + list(t_shape)) for t_shape in [x_shape, y_shape]]
|
||||
shape_ret = tuple(max(sx, sy) for sx,sy in zip(x_shape, y_shape))
|
||||
return input.reshape(x_shape).expand(shape_ret)
|
||||
return x.reshape(x_shape).expand(shape_ret)
|
||||
|
||||
# **************** Complex Ops ****************
|
||||
|
||||
@@ -182,9 +182,8 @@ def BatchNormalization(X: Tensor, scale, B, input_mean, input_var, epsilon=1e-05
|
||||
running_var = input_var * momentum + current_var * (1 - momentum)
|
||||
|
||||
return _batchnorm(X, scale, B, current_mean, current_invstd), running_mean, running_var
|
||||
else:
|
||||
invstd = (input_var + epsilon)**-0.5
|
||||
return _batchnorm(X, scale, B, input_mean, invstd)
|
||||
invstd = (input_var + epsilon)**-0.5
|
||||
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)))
|
||||
@@ -233,10 +232,10 @@ def _auto_pad(X: Tensor, auto_pad, strides, kernel_shape, dilations):
|
||||
pad_shape = [(math.ceil(sh/st)-1)*st+((ks-1)*di+1)-sh for sh, st, ks, di in zip(X.shape[-len(kernel_shape):], strides, kernel_shape, dilations)]
|
||||
pad_shape = flatten([[sh//2, sh-sh//2] for sh in pad_shape])
|
||||
return pad_shape[::2] + pad_shape[1::2] if auto_pad == "SAME_UPPER" else pad_shape[1::2] + pad_shape[::2]
|
||||
else: raise NotImplementedError(f"auto_pad={auto_pad} not implemented")
|
||||
raise NotImplementedError(f"auto_pad={auto_pad} not implemented")
|
||||
|
||||
def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Tensor=None, axes: Tensor=None, mode="constant", value: float=0.):
|
||||
constant_value = value if constant_value is None else float(safe_numpy(constant_value)[0])
|
||||
constant_value = value if constant_value is None else float(safe_numpy(constant_value))
|
||||
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
|
||||
@@ -247,7 +246,7 @@ def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Tensor=
|
||||
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":
|
||||
if mode == "reflect":
|
||||
for i,s in enumerate(x.shape):
|
||||
if pads[i] == (0,0): continue
|
||||
elif pads[i][0] and not pads[i][1]:
|
||||
@@ -261,7 +260,7 @@ def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Tensor=
|
||||
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":
|
||||
if mode == "edge":
|
||||
for i,s in enumerate(x.shape):
|
||||
if pads[i] == (0,0): continue
|
||||
elif pads[i][0] and not pads[i][1]:
|
||||
@@ -275,7 +274,7 @@ def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Tensor=
|
||||
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":
|
||||
if mode == "constant":
|
||||
return _padding(x, seq_pads, axes=seq_axes, constant_value=constant_value)
|
||||
|
||||
def AveragePool(X: Tensor, kernel_shape, auto_pad="NOTSET", ceil_mode=0, count_include_pad=0, dilations=1, pads=None, strides=1):
|
||||
@@ -283,9 +282,8 @@ def AveragePool(X: Tensor, kernel_shape, auto_pad="NOTSET", ceil_mode=0, count_i
|
||||
ret = _padding(X, pads, auto_pad, axes=pixel_axes, strides=strides, kernel_shape=kernel_shape, dilations=dilations, ceil_mode=ceil_mode).avg_pool2d(kernel_shape, stride=strides, dilation=dilations)
|
||||
if count_include_pad:
|
||||
return ret
|
||||
else:
|
||||
div = _padding(Tensor.ones(*X.shape), pads, auto_pad, axes=pixel_axes, strides=strides, kernel_shape=kernel_shape, dilations=dilations, ceil_mode=ceil_mode).avg_pool2d(kernel_shape, stride=strides, dilation=dilations)
|
||||
return ret / div
|
||||
div = _padding(Tensor.ones(*X.shape), pads, auto_pad, axes=pixel_axes, strides=strides, kernel_shape=kernel_shape, dilations=dilations, ceil_mode=ceil_mode).avg_pool2d(kernel_shape, stride=strides, dilation=dilations)
|
||||
return ret / div
|
||||
|
||||
def MaxPool(X: Tensor, kernel_shape, auto_pad="NOTSET", ceil_mode=0, dilations=1, pads=None, storage_order=0, strides=1):
|
||||
ret = _padding(X, pads, auto_pad, constant_value=float("-inf"), axes=tuple(range(len(X.shape)))[2:], strides=strides, kernel_shape=kernel_shape, dilations=dilations, ceil_mode=ceil_mode)
|
||||
@@ -343,21 +341,21 @@ def Dropout(data: Tensor, ratio=0.5, training_mode=False, seed=None):
|
||||
mask = Tensor((rng.random(data.shape) >= ratio), requires_grad=False, device=data.device)
|
||||
return data * mask * (1/(1.0 - ratio)), mask
|
||||
|
||||
def LRN(input: Tensor, size, alpha=1e-4, beta=0.75, bias=1.0):
|
||||
bs, c, iy, ix = input.shape
|
||||
return input / input.mul(input).reshape(bs,1,c,iy*ix).pad2d((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1).reshape(bs,c,iy,ix).mul(alpha).add(bias).pow(beta)
|
||||
def LRN(x: Tensor, size, alpha=1e-4, beta=0.75, bias=1.0):
|
||||
bs, c, iy, ix = x.shape
|
||||
return x / x.mul(x).reshape(bs,1,c,iy*ix).pad2d((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1).reshape(bs,c,iy,ix).mul(alpha).add(bias).pow(beta)
|
||||
|
||||
def MeanVarianceNormalization(input: Tensor, axis=(0, 2, 3)):
|
||||
data_mean = input.mean(axis=axis, keepdim=True)
|
||||
std = ((input**2).mean(axis=axis, keepdim=True) - data_mean**2).sqrt()
|
||||
return (input - data_mean) / (std + 1e-9)
|
||||
def MeanVarianceNormalization(x: Tensor, axis=(0, 2, 3)):
|
||||
data_mean = x.mean(axis=axis, keepdim=True)
|
||||
std = ((x**2).mean(axis=axis, keepdim=True) - data_mean**2).sqrt()
|
||||
return (x - data_mean) / (std + 1e-9)
|
||||
|
||||
def NegativeLogLikelihoodLoss(input: Tensor, target: Tensor, weight=None, ignore_index=None, reduction="mean"):
|
||||
def NegativeLogLikelihoodLoss(x: Tensor, target: Tensor, weight=None, ignore_index=None, reduction="mean"):
|
||||
target = target.cast(dtypes.float32)
|
||||
N, C, i_shape = input.shape[0], input.shape[1], input.shape
|
||||
N, C, i_shape = x.shape[0], x.shape[1], x.shape
|
||||
t_shape = target.shape
|
||||
if len(input.shape) != 3:
|
||||
input = input.reshape((N, C, -1))
|
||||
if len(x.shape) != 3:
|
||||
x = x.reshape((N, C, -1))
|
||||
target = target.reshape((N, -1))
|
||||
if weight is not None:
|
||||
mask = target.unsqueeze(-1) == Tensor.arange(C).repeat((N, 1, 1))
|
||||
@@ -365,14 +363,14 @@ def NegativeLogLikelihoodLoss(input: Tensor, target: Tensor, weight=None, ignore
|
||||
if ignore_index is not None:
|
||||
cond = target == ignore_index
|
||||
weight = cond.where(0, weight) if weight is not None else cond.where(Tensor.zeros(*target.shape), 1)
|
||||
mask = target[:, None, :] == Tensor.arange(C).reshape([1, C] + [1]*(len(input.shape) -2))
|
||||
loss = (-mask * input).sum(axis=1) * (1 if weight is None else weight)
|
||||
mask = target[:, None, :] == Tensor.arange(C).reshape([1, C] + [1]*(len(x.shape) -2))
|
||||
loss = (-mask * x).sum(axis=1) * (1 if weight is None else weight)
|
||||
if reduction == "mean": return loss.mean() if weight is None else loss.sum() / weight.sum()
|
||||
elif reduction == "sum": return loss.sum()
|
||||
if reduction == "sum": return loss.sum()
|
||||
return loss.reshape(t_shape) if len(i_shape) != 3 else loss
|
||||
|
||||
def SoftmaxCrossEntropyLoss(scores: Tensor, labels: Tensor, weights=None, ignore_index=None, reduction="mean"):
|
||||
N, C, *s_dimensions = scores.shape
|
||||
_N, C, *s_dimensions = scores.shape
|
||||
if ignore_index is not None: labels = (labels == ignore_index).where(C+1, labels).cast(dtypes.int32)
|
||||
mask = labels.unsqueeze(1) == Tensor.arange(C).reshape(1, C, *[1]*len(s_dimensions))
|
||||
y = scores.log_softmax(axis=1)
|
||||
@@ -382,30 +380,30 @@ def SoftmaxCrossEntropyLoss(scores: Tensor, labels: Tensor, weights=None, ignore
|
||||
elif reduction == "sum": loss = loss.sum()
|
||||
return loss, y
|
||||
|
||||
def ArrayFeatureExtractor(input: Tensor, indices: Tensor):
|
||||
return input.__getitem__(tuple([slice(None) if i != (input.ndim-1) else indices for i in range(input.ndim)]))
|
||||
def Gather(input: Tensor, indices: Tensor, axis=0):
|
||||
def ArrayFeatureExtractor(x: Tensor, indices: Tensor):
|
||||
return x.__getitem__(tuple([slice(None) if i != (x.ndim-1) else indices for i in range(x.ndim)]))
|
||||
def Gather(x: Tensor, indices: Tensor, 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:]
|
||||
x_sh = list(x.shape)
|
||||
ret_shape = x_sh[:axis] + list(indices.shape) + x_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)]))
|
||||
indices = [int(safe_numpy(indices))] if indices.shape == () else [x_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(x_sh)] for i in indices]
|
||||
return x.shrink(arg=tuple(args[0])).cat(*[x.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
|
||||
# NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
|
||||
return x.__getitem__(tuple([slice(None) if i != axis else indices for i in range(x.ndim)]))
|
||||
|
||||
def GatherElements(input: Tensor, indices: Tensor, axis):
|
||||
indices = (indices < 0).where(input.shape[axis], 0) + indices
|
||||
return input.gather(indices, axis)
|
||||
def GatherElements(x: Tensor, indices: Tensor, axis):
|
||||
indices = (indices < 0).where(x.shape[axis], 0) + indices
|
||||
return x.gather(indices, axis)
|
||||
|
||||
def _round(x:Tensor, n:float, equidistant_case = "round_down") -> Tensor:
|
||||
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":
|
||||
if equidistant_case == "round_up": return (x >= b).where(b+1-n, b-n)
|
||||
if equidistant_case == "round_to_even":
|
||||
def _and(cond1, cond2): return ((cond1 + cond2) == 2).where(1, 0)
|
||||
x_ceil_fraction = x.ceil()/2
|
||||
cond_ceil_even = x_ceil_fraction.ceil() == x_ceil_fraction
|
||||
@@ -443,8 +441,8 @@ def Resize(X:Tensor, roi=None, scales=None, sizes=None, antialias=0, axes=None,
|
||||
x_out = (x_out + 0.5)/scales_[-1] - 0.5 if output_shape[-1] > 1 else Tensor([0])
|
||||
y_out = (y_out + 0.5)/scales_[-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)])
|
||||
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)
|
||||
@@ -489,7 +487,7 @@ def Resize(X:Tensor, roi=None, scales=None, sizes=None, antialias=0, axes=None,
|
||||
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":
|
||||
if mode == "linear":
|
||||
x_out, y_out = _coordinate_transformation(x_out, y_out, output_shape_, scales, roi)
|
||||
ret = []
|
||||
for y in safe_numpy(y_out):
|
||||
@@ -509,18 +507,18 @@ def Resize(X:Tensor, roi=None, scales=None, sizes=None, antialias=0, axes=None,
|
||||
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")
|
||||
if mode == "cubic":
|
||||
raise NotImplementedError("cubic interpolation is not implemented")
|
||||
|
||||
def CenterCropPad(input: Tensor, shape: Tensor, axes=None):
|
||||
if not axes: axes = list(range(input.ndim))
|
||||
shrink_arg = [(0,i) for i in input.shape]
|
||||
pad_arg = [(0,0)] * input.ndim
|
||||
def CenterCropPad(t: Tensor, shape: Tensor, axes=None):
|
||||
if not axes: axes = list(range(t.ndim))
|
||||
shrink_arg = [(0,i) for i in t.shape]
|
||||
pad_arg = [(0,0)] * t.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)
|
||||
return input.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
|
||||
if s < t.shape[x]: shrink_arg[x] = (t.shape[x]//2 - s//2, t.shape[x]//2 + s//2) if s%2 == 0 else (t.shape[x]//2 - s//2 - 1, t.shape[x]//2 + s//2)
|
||||
elif s > t.shape[x]: pad_arg[x] = ((s - t.shape[x])//2, (s - t.shape[x])//2) if (s - t.shape[x])% 2 == 0 else ((s - t.shape[x])//2, (s - t.shape[x])//2 + 1)
|
||||
return t.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
|
||||
|
||||
def OneHot(indices: Tensor, depth: Tensor, values: Tensor, axis=-1):
|
||||
depth = int(safe_numpy(depth).item())
|
||||
@@ -558,11 +556,11 @@ def EyeLike(x: Tensor, dtype=None, k=0):
|
||||
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)
|
||||
if shape[0] == shape[1]:
|
||||
return Tensor.eye(dim=dim, dtype=dtype)
|
||||
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)
|
||||
|
||||
def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)
|
||||
|
||||
@@ -593,14 +591,13 @@ def ImageDecoder(encoded_stream: Tensor, pixel_format="RGB"):
|
||||
img = PIL.Image.open(io.BytesIO(safe_numpy(encoded_stream).tobytes()))
|
||||
if pixel_format == "BGR":
|
||||
return Tensor(np.array(img))[:, :, ::-1]
|
||||
elif pixel_format == "RGB":
|
||||
if pixel_format == "RGB":
|
||||
return Tensor(np.array(img))
|
||||
elif pixel_format == "Grayscale":
|
||||
if pixel_format == "Grayscale":
|
||||
img = img.convert("L")
|
||||
decoded = Tensor(np.array(img))
|
||||
return decoded.unsqueeze(-1) # (H, W) to (H, W, 1)
|
||||
else:
|
||||
raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
|
||||
raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
|
||||
|
||||
def AffineGrid(theta: Tensor, size: Tensor, align_corners=0):
|
||||
_, _, *data_sz = safe_numpy(size).tolist()
|
||||
@@ -619,20 +616,19 @@ def AffineGrid(theta: Tensor, size: Tensor, align_corners=0):
|
||||
assert dim_homo == 3
|
||||
original_grid = original_grid.reshape(H*W, dim_homo).transpose()
|
||||
return theta.matmul(original_grid).permute(0,2,1).reshape(N, H, W, dim_2d)
|
||||
else:
|
||||
assert original_grid.ndim == 4
|
||||
N, dim_3d, dim_homo = theta.shape
|
||||
assert dim_3d == 3 and dim_homo == 4
|
||||
D, H, W, dim_homo = original_grid.shape
|
||||
assert dim_homo == 4
|
||||
original_grid = original_grid.reshape(D*H*W, dim_homo).transpose()
|
||||
return theta.matmul(original_grid).permute(0,2,1).reshape(N, D, H, W, dim_3d)
|
||||
assert original_grid.ndim == 4
|
||||
N, dim_3d, dim_homo = theta.shape
|
||||
assert dim_3d == 3 and dim_homo == 4
|
||||
D, H, W, dim_homo = original_grid.shape
|
||||
assert dim_homo == 4
|
||||
original_grid = original_grid.reshape(D*H*W, dim_homo).transpose()
|
||||
return theta.matmul(original_grid).permute(0,2,1).reshape(N, D, H, W, dim_3d)
|
||||
|
||||
# **************** com.microsoft Ops ****************
|
||||
|
||||
def SkipLayerNormalization(input:Tensor, skip:Tensor, gamma, beta:Optional[Tensor]=None, bias:Optional[Tensor]=None, epsilon=None):
|
||||
def SkipLayerNormalization(x:Tensor, skip:Tensor, gamma, beta:Optional[Tensor]=None, bias:Optional[Tensor]=None, epsilon=None):
|
||||
if epsilon is None: epsilon=1e-12
|
||||
x = input + skip + bias
|
||||
x = x + skip + bias
|
||||
return x.layernorm(eps=epsilon) * gamma + beta, None, None, x
|
||||
|
||||
def FastGelu(x:Tensor, bias:Optional[Tensor]=None):
|
||||
@@ -645,7 +641,7 @@ def EmbedLayerNormalization(input_ids: Tensor, segment_ids:Optional[Tensor]=None
|
||||
assert (mask is None) is (mask_index_type is None)
|
||||
assert mask is None, "functionality not supported yet" # TODO
|
||||
input_shape = input_ids.shape
|
||||
bsz, seq_length = input_shape[0], input_shape[1]
|
||||
_bsz, seq_length = input_shape[0], input_shape[1]
|
||||
compute_seg_emb = (segment_embedding is not None and segment_ids is not None)
|
||||
vocab_size, max_position_embeddings, type_vocab_size = word_embedding.shape[0], position_embedding.shape[0], (segment_embedding.shape[0] if compute_seg_emb else None)
|
||||
|
||||
@@ -664,7 +660,7 @@ def EmbedLayerNormalization(input_ids: Tensor, segment_ids:Optional[Tensor]=None
|
||||
out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
|
||||
return out, None, embedding_sum
|
||||
|
||||
def Attention(input:Tensor, weights, bias:Optional[Tensor]=None, mask_index:Optional[Tensor]=None, past:Optional[Tensor]=None, relative_position_bias:Optional[Tensor]=None, past_sequence_length:Optional[Tensor]=None, do_rotary=None, mask_filter_value=None, num_heads=None, past_present_share_buffer=None, qkv_hidden_sizes=None, scale=None, unidirectional=None):
|
||||
def Attention(x:Tensor, weights, bias:Optional[Tensor]=None, mask_index:Optional[Tensor]=None, past:Optional[Tensor]=None, relative_position_bias:Optional[Tensor]=None, past_sequence_length:Optional[Tensor]=None, do_rotary=None, mask_filter_value=None, num_heads=None, past_present_share_buffer=None, qkv_hidden_sizes=None, scale=None, unidirectional=None):
|
||||
# https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.Attention
|
||||
assert num_heads is not None # required
|
||||
assert (qkv_hidden_sizes is None and past is not None) or (qkv_hidden_sizes is not None)
|
||||
@@ -673,12 +669,12 @@ def Attention(input:Tensor, weights, bias:Optional[Tensor]=None, mask_index:Opti
|
||||
|
||||
if unidirectional: # gpt-style
|
||||
assert hidden_size == v_hidden_size
|
||||
xqkv = input.linear(weights, bias)
|
||||
xqkv = x.linear(weights, bias)
|
||||
xq, xk, xv = [xqkv.slice([None, None, (i*hidden_size, (i+1)*hidden_size)]) for i in range(3)]
|
||||
else: # bert-style
|
||||
wq, wk, wv = weights[:,:hidden_size], weights[:,hidden_size:hidden_size+v_hidden_size], weights[:,hidden_size+v_hidden_size:]
|
||||
bq, bk, bv = (bias[:hidden_size], bias[hidden_size:hidden_size+v_hidden_size], bias[hidden_size+v_hidden_size]) if bias is not None else None
|
||||
xq, xk, xv = [input.linear(w, b) for w, b in zip((wq, wk, wv), (bq, bk, bv))]
|
||||
xq, xk, xv = [x.linear(w, b) for w, b in zip((wq, wk, wv), (bq, bk, bv))]
|
||||
xq, xk, xv = [x.reshape(x.shape[0], x.shape[1], num_heads, -1).transpose(1, 2) for x in (xq, xk, xv)]
|
||||
|
||||
if past is not None:
|
||||
@@ -703,11 +699,10 @@ def Attention(input:Tensor, weights, bias:Optional[Tensor]=None, mask_index:Opti
|
||||
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]
|
||||
T, R = safe_numpy(T), safe_numpy(R)
|
||||
r = R / (1 + T * decay_factor)
|
||||
ret = []
|
||||
for input in grouped_inputs:
|
||||
X, G, H = input
|
||||
for X, G, H in grouped_inputs:
|
||||
X.grad = norm_coefficient * X + G
|
||||
X.grad.requires_grad, H.requires_grad = False, False # TODO manually turning off requires_grad, see TODO under (domain == "ai.onnx.preview.training") in onnx.py
|
||||
H.assign(H.detach() + X.grad * X.grad).realize()
|
||||
@@ -720,11 +715,10 @@ def Adagrad(R, T, *inputs, decay_factor=0.0, epsilon=0.0, norm_coefficient=0.0):
|
||||
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]
|
||||
T, R = safe_numpy(T), safe_numpy(R)
|
||||
beta_adjusted = beta if T > 0 else 1
|
||||
ret = []
|
||||
for input in grouped_inputs:
|
||||
X, G, V = input
|
||||
for X, G, V in grouped_inputs:
|
||||
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()
|
||||
@@ -738,10 +732,9 @@ def Momentum(R, T, *inputs, alpha, beta, mode, norm_coefficient):
|
||||
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]
|
||||
T, R = safe_numpy(T), safe_numpy(R)
|
||||
ret = []
|
||||
for input in grouped_inputs:
|
||||
X, G, V, H = input
|
||||
for X, G, V, H in grouped_inputs:
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user