From e263c0c628ffde9e853341d938fa90cb89bea794 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 24 Feb 2023 09:22:58 -0800 Subject: [PATCH] onnx : another model test is passing --- extra/onnx.py | 15 +++------------ extra/onnx_ops.py | 28 ++++++++++++++++++++++++++++ test/external_test_onnx_backend.py | 19 +++++++++++-------- 3 files changed, 42 insertions(+), 20 deletions(-) diff --git a/extra/onnx.py b/extra/onnx.py index de5bc5ef1f..207212914c 100644 --- a/extra/onnx.py +++ b/extra/onnx.py @@ -9,6 +9,7 @@ from onnx.helper import tensor_dtype_to_np_dtype # global numpy cache for parameters numpy_cache = {} def safe_numpy(t): + if not isinstance(t, Tensor): return t global numpy_cache if t not in numpy_cache: if DEBUG >= 1: @@ -126,10 +127,6 @@ def get_run_onnx(onnx_model): 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 == "Conv": - x,w,b = inp if len(inp) == 3 else (inp[0], inp[1], None) - assert 'dilations' not in opt or opt['dilations'] == (1,1) - ret = x.conv2d(w, b, stride=opt['strides'], groups=opt.get('group', 1), padding=(opt['pads'][0], opt['pads'][2], opt['pads'][1], opt['pads'][3]) if 'pads' in opt else 0) elif n.op_type in ["Sum"]: ret = functools.reduce(Tensor.__add__, inp) elif n.op_type in ["Add", "Sub", "Mul"]: @@ -150,12 +147,6 @@ def get_run_onnx(onnx_model): intermediate_tensors[o] = inp[0].slice(arg=arg) i = i+s continue - elif n.op_type == "AveragePool": - ret = inp[0].pad2d((opt['pads'][0], opt['pads'][2], opt['pads'][1], opt['pads'][3])) if 'pads' in opt else inp[0] - ret = ret.avg_pool2d(opt['kernel_shape'], opt.get('strides', [1]*len(opt['kernel_shape']))) - elif n.op_type == "MaxPool": - ret = inp[0].pad2d((opt['pads'][0], opt['pads'][2], opt['pads'][1], opt['pads'][3])) if 'pads' in opt else inp[0] - ret = ret.max_pool2d(opt['kernel_shape'], opt.get('strides', [1]*len(opt['kernel_shape']))) elif n.op_type == "Slice": assert onnx_model.opset_import[0].version == 10 arg = [(0,x) for x in inp[0].shape] @@ -172,9 +163,9 @@ def get_run_onnx(onnx_model): print("UNSUPPORTED", n.op_type, n.input, n.output) raise Exception(f"op_type {n.op_type} not supported") if not isinstance(ret, tuple): ret = (ret, ) - assert len(n.output) == len(ret), f"output size must be {len(ret)}, it's {n.output}" + assert len(n.output) <= len(ret), f"expected output size must be less than {len(ret)}, it's {n.output}" if debug: print([x.shape for x in ret]) - for i,r in enumerate(ret): intermediate_tensors[n.output[i]] = r + for i in range(len(n.output)): intermediate_tensors[n.output[i]] = ret[i] #print(ret.numpy().mean()) if num == ONNXLIMIT: output_tensor_names = n.output diff --git a/extra/onnx_ops.py b/extra/onnx_ops.py index 6fcfddc8e5..851ed21e8c 100644 --- a/extra/onnx_ops.py +++ b/extra/onnx_ops.py @@ -1,4 +1,6 @@ +from tinygrad.tensor import Tensor from extra.onnx import safe_numpy +import numpy as np def Unsqueeze(data, axes): axes = [len(data.shape) + int(x) if x < 0 else int(x) for x in safe_numpy(axes)] @@ -32,3 +34,29 @@ def BatchNormalization(X, scale, B, input_mean, input_var, epsilon=1e-05, moment else: invstd = (input_var + epsilon)**-0.5 return X.batchnorm(scale, B, input_mean, invstd) + +def _padding(pads=None, auto_pad="NOTSET"): + assert auto_pad == "NOTSET" # TODO: write this + return (pads[1], pads[3], pads[0], pads[2]) if pads is not None else (0,0,0,0) + +def AveragePool(X, kernel_shape, auto_pad="NOTSET", ceil_mode=0, count_include_pad=0, dilations=1, pads=None, strides=1): + # TODO: the padding shouldn't be counted in the average! this is causing a test failure + assert ceil_mode == 0 and count_include_pad == 0 and dilations == 1 + return X.pad2d(_padding(pads, auto_pad)).avg_pool2d(kernel_shape, stride=strides) + +def MaxPool(X, kernel_shape, auto_pad="NOTSET", ceil_mode=0, dilations=1, pads=None, storage_order=0, strides=1): + # TODO: the padding should be infinity, not 0! + assert ceil_mode == 0 and storage_order == 0 and dilations == 1 + return X.pad2d(_padding(pads, auto_pad)).max_pool2d(kernel_shape, stride=strides) + +def Conv(X, W, B=None, auto_pad="NOTSET", dilations=1, group=1, kernel_shape=None, pads=None, strides=1): + return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations, padding=_padding(pads, auto_pad)) + +# TODO: copied from tensor.py +def Dropout(data, ratio=0.5, training_mode=False, seed=None): + # TODO: mask should be a boolean tensor + if not training_mode: return data, Tensor.ones(*data.shape) # if mask is requested as output it will contain all ones. + if seed is not None: Tensor.manual_seed(seed) + _mask : np.ndarray = np.asarray(Tensor._rng.binomial(1, 1.0-ratio, size=data.shape), dtype=data.dtype) + mask = Tensor(_mask, requires_grad=False, device=data.device) + return data * mask * (1/(1.0 - ratio)), mask \ No newline at end of file diff --git a/test/external_test_onnx_backend.py b/test/external_test_onnx_backend.py index 8954db5b4f..3ab7980fdc 100644 --- a/test/external_test_onnx_backend.py +++ b/test/external_test_onnx_backend.py @@ -45,13 +45,20 @@ backend_test.include('test_unsqueeze_*') backend_test.include('test_gemm_*') backend_test.include('test_batchnorm_*') +# almost passing node tests +#backend_test.include('test_conv_.*') +#backend_test.include('test_dropout_*') + +# failing for real reasons +#backend_test.include('test_averagepool_2d_*') +#backend_test.include('test_maxpool_2d_*') + """ backend_test.include('test_sum_*') backend_test.include('test_transpose_*') backend_test.include('test_tanh_*') # should be passing (good place to start!) -backend_test.include('test_conv_.*') backend_test.include('test_reshape_*') backend_test.include('test_flatten_*') backend_test.include('test_expand_*') @@ -74,17 +81,17 @@ backend_test.include('test_clip_*') #backend_test.include('test_slice_*') #backend_test.include('test_lrn_*') #backend_test.include('test_batchnorm_*') -#backend_test.include('test_maxpool_*') -#backend_test.include('test_averagepool_*') -""" # working big model tests backend_test.include('test_resnet50') backend_test.include('test_densenet121') +backend_test.include('test_vgg19') +""" # wrong big model tests backend_test.include('test_shufflenet') backend_test.include('test_inception_v2') +backend_test.include('test_squeezenet') """ """ @@ -92,10 +99,6 @@ backend_test.include('test_inception_v2') backend_test.include('test_bvlc_alexnet') backend_test.include('test_inception_v1') backend_test.include('test_zfnet512') - -# unsupported big model tests : Dropout -backend_test.include('test_squeezenet') -backend_test.include('test_vgg19') """ globals().update(backend_test.enable_report().test_cases)