diff --git a/extra/onnx.py b/extra/onnx.py index bf31e9355b..8f838b00a2 100644 --- a/extra/onnx.py +++ b/extra/onnx.py @@ -195,6 +195,7 @@ def get_onnx_ops(): return [pads[i]-pads[i]//2 for i in range(len(pads))] + [pads[i]//2 for i in range(len(pads))] def _resolve_pool_pads(x:Tensor, p_, k_, d_, s_, auto_pad:AUTO_PAD_OPTIONS): + if auto_pad == "VALID": return [0]*(len(k_)*2) i_, (s_,d_,p_) = x.shape[-len(k_):], (make_tuple(x, len(k_)*2) for x in (s_, d_, p_)) if auto_pad == "NOTSET": return _onnx_pads_to_tiny_pads(p_ if len(p_)==len(k_)*2 else p_*2) o_ = [((i - (1 if auto_pad in ("SAME_UPPER", "SAME_LOWER") else k)) // s + 1) for i,k,s in zip(i_, k_, s_)] @@ -673,7 +674,8 @@ def get_onnx_ops(): 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 = [_cached_to_python_const(indices)] if indices.shape == () else [x_sh[axis]+x if x<0 else x for x in _cached_to_python_const(indices)] + indices = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices) + indices = [x_sh[axis]+x if x<0 else x for x in indices] args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x_sh)] for i in indices] # type: ignore 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 diff --git a/test/external/external_test_onnx_ops.py b/test/external/external_test_onnx_ops.py index 1442a0a565..7e846695f9 100644 --- a/test/external/external_test_onnx_ops.py +++ b/test/external/external_test_onnx_ops.py @@ -27,6 +27,27 @@ class TestMainOnnxOps(TestOnnxOps): outputs = ["out"] self.helper_test_single_op("Reshape", inputs, attributes, outputs) + def test_conv(self): + # test VALID auto_pad + inputs = { + "x": np.random.randn(1, 3, 384, 384).astype(np.float32), + "w": np.random.randn(1152, 3, 14, 14).astype(np.float32), + "b": np.random.randn(1152).astype(np.float32) + } + attributes = {'auto_pad': 'VALID', 'dilations': (1, 1), 'group': 1, 'kernel_shape': (14, 14), 'strides': (14, 14)} + outputs = ["y"] + self.helper_test_single_op("Conv", inputs, attributes, outputs, atol=1e-4) + + def test_gather(self): + # test const negative indices + inputs = { + "input": np.random.randn(1, 3, 3).astype(np.float32), + "indices": np.array(-2, dtype=np.int64), + } + attributes = {'axis': 1} + outputs = ["y"] + self.helper_test_single_op("Gather", inputs, attributes, outputs) + def test_quantize_linear(self): test_cases = [ {"test_case": "round_half_to_even", "qdtype": np.int8, "qzero_point": 0, "x": [-1.5, -0.5, 0.5, 1.5], "scale": 1.0},