From fb76dc30c73f099840fc87631d45f20798e6c27f Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Thu, 4 Jan 2024 19:06:03 +0800 Subject: [PATCH 01/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 21 ++++++++++++++++--- .../tests/function_libs/torch_lib/ops_test.py | 1 + .../function_libs/torch_lib/ops_test_data.py | 17 +++++++++++++++ 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index bb767071e7..14dd8422f6 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2391,12 +2391,27 @@ def aten_upsample_linear1d_backward( raise NotImplementedError() +@torch_op("aten::upsample_nearest1d", trace_only=True) def aten_upsample_nearest1d( - self: TensorType, output_size: INT64, scales: Optional[float] = None -) -> TensorType: + self: TReal, size: INT64, scale_factor: Optional[float] = None +) -> TReal: """upsample_nearest1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor""" + if size is not None: + result = _aten_upsample_nearest2d_onnx(self, size) + else: + result = _aten_upsample_nearest1d_scales_onnx(self, scale_factor) + return result - raise NotImplementedError() + +@torch_op("aten::upsample_nearest1d", private=True) +def _aten_upsample_nearest1d_scales_onnx( + self: TReal, output_size: INT64 +) -> TReal: + + self_shape = op.Shape(self) + batch_channel = self_shape[:2] + output_size = op.Concat(batch_channel, output_size, axis=0) + return op.Resize(self, None, None, output_size, mode="nearest", coordinate_transformation_mode="pytorch_half_pixel", nearest_mode="floor") def aten_upsample_nearest1d_backward( diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py index 9cae237c80..1a3e68fadb 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py @@ -199,6 +199,7 @@ def run_test_output_match( ), kwargs=repr(cpu_sample.kwargs), ): + #if i != 1: continue test_behavior, reason = _should_skip_xfail_test_sample(op.name, cpu_sample, dtype) with ops_test_common.normal_xfail_skip_test_behaviors(test_behavior, reason): diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index 7aeba0d14b..0d810dbc03 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -439,6 +439,17 @@ def _upsample_input_wrangler( return args, kwargs +def _upsample_nearest1d_input_wrangler( + args: list[Any], kwargs: dict[str, Any] +) -> tuple[list[Any], dict[str, Any]]: + # if "scale_factor" in kwargs: + # del kwargs["scale_factor"] + # if "size" in kwargs: + # kwargs["output_size"] = kwargs["size"] + # del kwargs["size"] + return args, kwargs + + def _unflatten_input_wrangler( args: list[Any], kwargs: dict[str, Any] ) -> tuple[list[Any], dict[str, Any]]: @@ -2127,6 +2138,12 @@ def _where_input_wrangler( nn_ops.aten_upsample_bicubic2d, trace_only=True, ), + TorchLibOpInfo( + "nn.functional.upsample_nearest1d", + nn_ops.aten_upsample_nearest1d, + input_wrangler=_upsample_nearest1d_input_wrangler, + trace_only=True, + ), TorchLibOpInfo( "nn.functional.upsample_nearest2d", nn_ops.aten_upsample_nearest2d, From 2a90007b995a2af4f7daa265e4aff839474fa005 Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Fri, 5 Jan 2024 17:03:10 +0800 Subject: [PATCH 02/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 17 +---------------- .../tests/function_libs/torch_lib/ops_test.py | 2 +- .../function_libs/torch_lib/ops_test_data.py | 4 ++++ 3 files changed, 6 insertions(+), 17 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index 14dd8422f6..ae603fa4cf 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2396,22 +2396,7 @@ def aten_upsample_nearest1d( self: TReal, size: INT64, scale_factor: Optional[float] = None ) -> TReal: """upsample_nearest1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor""" - if size is not None: - result = _aten_upsample_nearest2d_onnx(self, size) - else: - result = _aten_upsample_nearest1d_scales_onnx(self, scale_factor) - return result - - -@torch_op("aten::upsample_nearest1d", private=True) -def _aten_upsample_nearest1d_scales_onnx( - self: TReal, output_size: INT64 -) -> TReal: - - self_shape = op.Shape(self) - batch_channel = self_shape[:2] - output_size = op.Concat(batch_channel, output_size, axis=0) - return op.Resize(self, None, None, output_size, mode="nearest", coordinate_transformation_mode="pytorch_half_pixel", nearest_mode="floor") + return _aten_upsample_nearest2d_onnx(self, size) def aten_upsample_nearest1d_backward( diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py index 1a3e68fadb..e03cde40e0 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py @@ -199,7 +199,7 @@ def run_test_output_match( ), kwargs=repr(cpu_sample.kwargs), ): - #if i != 1: continue + #if i != 10: continue test_behavior, reason = _should_skip_xfail_test_sample(op.name, cpu_sample, dtype) with ops_test_common.normal_xfail_skip_test_behaviors(test_behavior, reason): diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index 0d810dbc03..85492790a6 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -2143,6 +2143,10 @@ def _where_input_wrangler( nn_ops.aten_upsample_nearest1d, input_wrangler=_upsample_nearest1d_input_wrangler, trace_only=True, + ).skip( + # size parameter must be existed + matcher=lambda sample: sample.kwargs.get("size", None) is None, + reason="aten_upsample_nearest1d takes size as input", ), TorchLibOpInfo( "nn.functional.upsample_nearest2d", From 0e9e65ee02ff07dbdbfd7a1755108c3318dbdca7 Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Fri, 5 Jan 2024 17:04:24 +0800 Subject: [PATCH 03/13] update --- onnxscript/tests/function_libs/torch_lib/ops_test.py | 1 - .../tests/function_libs/torch_lib/ops_test_data.py | 12 ------------ 2 files changed, 13 deletions(-) diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py index e03cde40e0..9cae237c80 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py @@ -199,7 +199,6 @@ def run_test_output_match( ), kwargs=repr(cpu_sample.kwargs), ): - #if i != 10: continue test_behavior, reason = _should_skip_xfail_test_sample(op.name, cpu_sample, dtype) with ops_test_common.normal_xfail_skip_test_behaviors(test_behavior, reason): diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index 85492790a6..edcadfb171 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -439,17 +439,6 @@ def _upsample_input_wrangler( return args, kwargs -def _upsample_nearest1d_input_wrangler( - args: list[Any], kwargs: dict[str, Any] -) -> tuple[list[Any], dict[str, Any]]: - # if "scale_factor" in kwargs: - # del kwargs["scale_factor"] - # if "size" in kwargs: - # kwargs["output_size"] = kwargs["size"] - # del kwargs["size"] - return args, kwargs - - def _unflatten_input_wrangler( args: list[Any], kwargs: dict[str, Any] ) -> tuple[list[Any], dict[str, Any]]: @@ -2141,7 +2130,6 @@ def _where_input_wrangler( TorchLibOpInfo( "nn.functional.upsample_nearest1d", nn_ops.aten_upsample_nearest1d, - input_wrangler=_upsample_nearest1d_input_wrangler, trace_only=True, ).skip( # size parameter must be existed From 48b994bb08a87489618f01cad18a897d4027e987 Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Thu, 11 Jan 2024 15:34:58 +0800 Subject: [PATCH 04/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 9 +++++---- onnxscript/tests/function_libs/torch_lib/ops_test.py | 1 + .../tests/function_libs/torch_lib/ops_test_data.py | 5 +++++ 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index ae603fa4cf..ddf62b313b 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2461,16 +2461,17 @@ def aten_upsample_nearest2d_backward( raise NotImplementedError() +@torch_op("aten::upsample_nearest3d", trace_only=True) def aten_upsample_nearest3d( - self: TensorType, - output_size: INT64, + self: TReal, + size: INT64, scales_d: Optional[float] = None, scales_h: Optional[float] = None, scales_w: Optional[float] = None, -) -> TensorType: +) -> TReal: """upsample_nearest3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor""" - raise NotImplementedError() + return op.Identity(self) def aten_upsample_nearest3d_backward( diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py index 9cae237c80..712f628a43 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py @@ -199,6 +199,7 @@ def run_test_output_match( ), kwargs=repr(cpu_sample.kwargs), ): + if i != 0: continue test_behavior, reason = _should_skip_xfail_test_sample(op.name, cpu_sample, dtype) with ops_test_common.normal_xfail_skip_test_behaviors(test_behavior, reason): diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index edcadfb171..974c42f9cd 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -2151,6 +2151,11 @@ def _where_input_wrangler( matcher=lambda sample: "scale_factor" in sample.kwargs, reason="fixme: the scale_factor tests", ), + TorchLibOpInfo( + "nn.functional.upsample_nearest3d", + nn_ops.aten_upsample_nearest3d, + trace_only=True, + ), TorchLibOpInfo("ones_like", core_ops.aten_ones_like, trace_only=True), TorchLibOpInfo( "roll", From edd16c87785866de544ade181fab98dd222e559a Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Mon, 15 Jan 2024 20:20:18 +0800 Subject: [PATCH 05/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 7 +++++-- onnxscript/tests/function_libs/torch_lib/ops_test_data.py | 4 ++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index ddf62b313b..3caa6cdc94 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2433,7 +2433,9 @@ def _aten_upsample_nearest2d_onnx( size: INT64, ) -> TReal: self_shape = op.Shape(self) - batch_channel = self_shape[:2] # type: ignore[index] + starts = op.Constant(value_ints=[0]) + ends = op.Constant(value_ints=[2]) + batch_channel = op.Slice(self_shape, starts, ends) output_size = op.Concat(batch_channel, size, axis=0) return op.Resize( @@ -2471,7 +2473,8 @@ def aten_upsample_nearest3d( ) -> TReal: """upsample_nearest3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor""" - return op.Identity(self) + return _aten_upsample_nearest2d_onnx(self, size) + def aten_upsample_nearest3d_backward( diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index 974c42f9cd..d4e29553ee 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -2155,6 +2155,10 @@ def _where_input_wrangler( "nn.functional.upsample_nearest3d", nn_ops.aten_upsample_nearest3d, trace_only=True, + ).skip( + # Shape should be [N, C, H, W] + matcher=lambda sample: len(sample.input.shape) != 2 + 3, + reason="only test on 2d inputs", ), TorchLibOpInfo("ones_like", core_ops.aten_ones_like, trace_only=True), TorchLibOpInfo( From 1fff4da06ba2ab6f3d2c510771a51825ec3a5fd9 Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Mon, 22 Jan 2024 18:57:23 +0800 Subject: [PATCH 06/13] update --- .../function_libs/torch_lib/extra_opinfo.py | 111 ++++++++++++++++++ .../tests/function_libs/torch_lib/ops_test.py | 2 +- .../function_libs/torch_lib/ops_test_data.py | 16 +-- 3 files changed, 116 insertions(+), 13 deletions(-) diff --git a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py index 0224b6cfa4..b8a91348a3 100644 --- a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py +++ b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py @@ -1460,6 +1460,103 @@ def shape(size, rank, with_batch_channel=True): ) +def sample_inputs_upsample_nearest1d(op_info, device, dtype, requires_grad, **kwargs): + del op_info + del kwargs + + N, C = 2, 3 + D = 4 + SS = 3 + L = 5 + + rank = 1 + + def shape(size, rank, with_batch_channel=True): + if with_batch_channel: + return tuple([N, C] + ([size] * rank)) + return tuple([size] * rank) + + make_arg = functools.partial( + torch_testing.make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + low=-1, + high=1, + ) + + yield opinfo_core.SampleInput(make_arg(shape(D, rank)), shape(SS, rank, False), True) + + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), shape(S, rank, False), + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + shape(L, rank, False), + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + None, # output_size + (1.7,), # scaler + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + None, # if this is None, the scalar must be list + 0.6, + ) + + +def sample_inputs_upsample_nearest3d(op_info, device, dtype, requires_grad, **kwargs): + del op_info + del kwargs + + N, C = 2, 3 + D = 4 + SS = 3 + L = 5 + + align_corners_options = (True, False) + rank = 3 + + def shape(size, rank, with_batch_channel=True): + if with_batch_channel: + return tuple([N, C] + ([size] * rank)) + return tuple([size] * rank) + + make_arg = functools.partial( + torch_testing.make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + low=-1, + high=1, + ) + + yield opinfo_core.SampleInput(make_arg(shape(D, rank)), shape(SS, rank, False), True) + + for align_corners in align_corners_options: + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), shape(S, rank, False), align_corners + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + shape(L, rank, False), + align_corners, + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + None, # output_size + align_corners, + (1.7, 1.7), # scaler + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + None, # if this is None, the scalar must be list + align_corners, + (0.6, 0.6), + ) + + class _TestParamsMaxPoolEmptyStrideBase: # Adapted from https://github.com/pytorch/pytorch/blob/d6d55f8590eab05d2536756fb4efcfb2d07eb81a/torch/testing/_internal/common_methods_invocations.py#L3203 def __init__(self): @@ -1932,6 +2029,20 @@ def __init__(self): sample_inputs_func=sample_inputs_upsample_bicubic2d, supports_out=False, ), + opinfo_core.OpInfo( + "ops.aten.upsample_nearest1d", + aten_name="upsample_nearest1d", + dtypes=common_dtype.floating_types_and(torch.bfloat16), + sample_inputs_func=sample_inputs_upsample_nearest1d, + supports_out=False, + ), + opinfo_core.OpInfo( + "ops.aten.upsample_nearest3d", + aten_name="upsample_nearest3d", + dtypes=common_dtype.floating_types_and(torch.bfloat16), + sample_inputs_func=sample_inputs_upsample_nearest3d, + supports_out=False, + ), opinfo_core.OpInfo( "nn.functional.max_pool1d_with_indices", aten_name="max_pool1d_with_indices", diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py index 712f628a43..2fa22e07d5 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py @@ -199,7 +199,7 @@ def run_test_output_match( ), kwargs=repr(cpu_sample.kwargs), ): - if i != 0: continue + if i != 3: continue test_behavior, reason = _should_skip_xfail_test_sample(op.name, cpu_sample, dtype) with ops_test_common.normal_xfail_skip_test_behaviors(test_behavior, reason): diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index d4e29553ee..2949775adc 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -2128,13 +2128,9 @@ def _where_input_wrangler( trace_only=True, ), TorchLibOpInfo( - "nn.functional.upsample_nearest1d", + "ops.aten.upsample_nearest1d", nn_ops.aten_upsample_nearest1d, trace_only=True, - ).skip( - # size parameter must be existed - matcher=lambda sample: sample.kwargs.get("size", None) is None, - reason="aten_upsample_nearest1d takes size as input", ), TorchLibOpInfo( "nn.functional.upsample_nearest2d", @@ -2152,13 +2148,9 @@ def _where_input_wrangler( reason="fixme: the scale_factor tests", ), TorchLibOpInfo( - "nn.functional.upsample_nearest3d", + "ops.aten.upsample_nearest3d", nn_ops.aten_upsample_nearest3d, trace_only=True, - ).skip( - # Shape should be [N, C, H, W] - matcher=lambda sample: len(sample.input.shape) != 2 + 3, - reason="only test on 2d inputs", ), TorchLibOpInfo("ones_like", core_ops.aten_ones_like, trace_only=True), TorchLibOpInfo( @@ -2384,9 +2376,9 @@ def _where_input_wrangler( OPS_DB, "nn.functional.upsample_nearest", ( - "nn.functional.upsample_nearest1d", + # "nn.functional.upsample_nearest1d", "nn.functional.upsample_nearest2d", - "nn.functional.upsample_nearest3d", + # "nn.functional.upsample_nearest3d", ), ) ops_test_common.duplicate_opinfo( From aba0961660f280ac224249f5fa91be1949dac07c Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Thu, 25 Jan 2024 17:35:43 +0800 Subject: [PATCH 07/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 32 ++++++++++++-- .../function_libs/torch_lib/extra_opinfo.py | 42 +++++++++---------- .../tests/function_libs/torch_lib/ops_test.py | 1 - 3 files changed, 47 insertions(+), 28 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index b56e389eb7..d4a2796488 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2401,7 +2401,28 @@ def aten_upsample_nearest1d( self: TReal, size: INT64, scale_factor: Optional[float] = None ) -> TReal: """upsample_nearest1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor""" - return _aten_upsample_nearest2d_onnx(self, size) + if size is not None: + return _aten_upsample_nearest_size_onnx(self, size) + else: + return _aten_upsample_nearest1d_scales(self, scale_factor) + + +@torch_op("aten::upsample_nearest1d", private=True) +def _aten_upsample_nearest1d_scales( + self: TReal, + scale_factors: TFloat, +) -> TReal: + scale_factors = op.Cast(scale_factors, to=FLOAT.dtype) + scale_factors = op.Concat(op.Constant(value_floats=[1.0, 1.0]), scale_factors, axis=0) + return op.Resize( + self, + None, + scale_factors, # format should be: [1.0, 1.0, scale_h, scale_w] + None, + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) def aten_upsample_nearest1d_backward( @@ -2429,11 +2450,11 @@ def aten_upsample_nearest2d( del scales_h del scales_w - return _aten_upsample_nearest2d_onnx(self, size) + return _aten_upsample_nearest_size_onnx(self, size) @torch_op("aten::upsample_nearest2d", private=True) -def _aten_upsample_nearest2d_onnx( +def _aten_upsample_nearest_size_onnx( self: TReal, size: INT64, ) -> TReal: @@ -2478,8 +2499,11 @@ def aten_upsample_nearest3d( ) -> TReal: """upsample_nearest3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor""" - return _aten_upsample_nearest2d_onnx(self, size) + del scales_h + del scales_w + del scales_d + return _aten_upsample_nearest_size_onnx(self, size) def aten_upsample_nearest3d_backward( diff --git a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py index 3b2390ad93..fb81e95ab4 100644 --- a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py +++ b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py @@ -1621,7 +1621,7 @@ def shape(size, rank, with_batch_channel=True): yield opinfo_core.SampleInput( make_arg(shape(D, rank)), None, # if this is None, the scalar must be list - 0.6, + (0.6,), ) @@ -1634,7 +1634,6 @@ def sample_inputs_upsample_nearest3d(op_info, device, dtype, requires_grad, **kw SS = 3 L = 5 - align_corners_options = (True, False) rank = 3 def shape(size, rank, with_batch_channel=True): @@ -1653,27 +1652,24 @@ def shape(size, rank, with_batch_channel=True): yield opinfo_core.SampleInput(make_arg(shape(D, rank)), shape(SS, rank, False), True) - for align_corners in align_corners_options: - yield opinfo_core.SampleInput( - make_arg(shape(D, rank)), shape(S, rank, False), align_corners - ) - yield opinfo_core.SampleInput( - make_arg(shape(D, rank)), - shape(L, rank, False), - align_corners, - ) - yield opinfo_core.SampleInput( - make_arg(shape(D, rank)), - None, # output_size - align_corners, - (1.7, 1.7), # scaler - ) - yield opinfo_core.SampleInput( - make_arg(shape(D, rank)), - None, # if this is None, the scalar must be list - align_corners, - (0.6, 0.6), - ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), shape(S, rank, False), + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + shape(L, rank, False), + ) + # ONNX don't support below cases: both output_size and scaler are not None + # yield opinfo_core.SampleInput( + # make_arg(shape(D, rank)), + # shape(L, rank, False), + # 1.7, # scaler + # ) + # yield opinfo_core.SampleInput( + # make_arg(shape(D, rank)), + # shape(L, rank, False), + # 0.6, + # ) class _TestParamsMaxPoolEmptyStrideBase: diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py index 2fa22e07d5..9cae237c80 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py @@ -199,7 +199,6 @@ def run_test_output_match( ), kwargs=repr(cpu_sample.kwargs), ): - if i != 3: continue test_behavior, reason = _should_skip_xfail_test_sample(op.name, cpu_sample, dtype) with ops_test_common.normal_xfail_skip_test_behaviors(test_behavior, reason): From 18630cd576afd3214b1a57b1d6bb9323632f893f Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Thu, 25 Jan 2024 17:59:13 +0800 Subject: [PATCH 08/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 33 +++---------------- .../function_libs/torch_lib/extra_opinfo.py | 6 ++-- 2 files changed, 9 insertions(+), 30 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index d4a2796488..d1a1c3e039 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2207,6 +2207,7 @@ def _aten_upsample_output_size( output_size: INT64, mode: str, coordinate_transformation_mode: str, + nearest_mode: str = "floor", ) -> TReal: self_shape = op.Shape(self) starts = op.Constant(value_ints=[0]) @@ -2220,7 +2221,7 @@ def _aten_upsample_output_size( output_size, mode=mode, coordinate_transformation_mode=coordinate_transformation_mode, - nearest_mode="floor", + nearest_mode=nearest_mode, ) @@ -2402,7 +2403,7 @@ def aten_upsample_nearest1d( ) -> TReal: """upsample_nearest1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor""" if size is not None: - return _aten_upsample_nearest_size_onnx(self, size) + return _aten_upsample_output_size(self, size, "nearest", "asymmetric", "floor") else: return _aten_upsample_nearest1d_scales(self, scale_factor) @@ -2450,31 +2451,7 @@ def aten_upsample_nearest2d( del scales_h del scales_w - return _aten_upsample_nearest_size_onnx(self, size) - - -@torch_op("aten::upsample_nearest2d", private=True) -def _aten_upsample_nearest_size_onnx( - self: TReal, - size: INT64, -) -> TReal: - self_shape = op.Shape(self) - starts = op.Constant(value_ints=[0]) - ends = op.Constant(value_ints=[2]) - batch_channel = op.Slice(self_shape, starts, ends) - output_size = op.Concat(batch_channel, size, axis=0) - - return op.Resize( - self, - None, - None, - output_size, - mode="nearest", - # NOTE(justinchuby): Both asymmetric and pytorch_half_pixel pass the test - # I used asymmetric because it aligns with the torch.onnx exporter - coordinate_transformation_mode="asymmetric", - nearest_mode="floor", - ) + return _aten_upsample_output_size(self, size, "nearest", "asymmetric", "floor") def aten_upsample_nearest2d_backward( @@ -2503,7 +2480,7 @@ def aten_upsample_nearest3d( del scales_w del scales_d - return _aten_upsample_nearest_size_onnx(self, size) + return _aten_upsample_output_size(self, size, "nearest", "asymmetric", "floor") def aten_upsample_nearest3d_backward( diff --git a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py index fb81e95ab4..96d38905cd 100644 --- a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py +++ b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py @@ -1607,7 +1607,8 @@ def shape(size, rank, with_batch_channel=True): yield opinfo_core.SampleInput(make_arg(shape(D, rank)), shape(SS, rank, False), True) yield opinfo_core.SampleInput( - make_arg(shape(D, rank)), shape(S, rank, False), + make_arg(shape(D, rank)), + shape(S, rank, False), ) yield opinfo_core.SampleInput( make_arg(shape(D, rank)), @@ -1653,7 +1654,8 @@ def shape(size, rank, with_batch_channel=True): yield opinfo_core.SampleInput(make_arg(shape(D, rank)), shape(SS, rank, False), True) yield opinfo_core.SampleInput( - make_arg(shape(D, rank)), shape(S, rank, False), + make_arg(shape(D, rank)), + shape(S, rank, False), ) yield opinfo_core.SampleInput( make_arg(shape(D, rank)), From 6d1e47c777b7b745c5b431876d799a6bbddaae1e Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Fri, 26 Jan 2024 14:13:09 +0800 Subject: [PATCH 09/13] Update nn.py --- onnxscript/function_libs/torch_lib/ops/nn.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index d1a1c3e039..248e7c26cf 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2201,7 +2201,16 @@ def _get_upsample_align_corners_mode(align_corners: bool) -> str: return "align_corners" if align_corners else "pytorch_half_pixel" -@torch_op(("aten::upsample_bicubic2d", "aten::upsample_bilinear2d"), private=True) +@torch_op( + ( + "aten::upsample_bicubic2d", + "aten::upsample_bilinear2d", + "aten::upsample_nearest1d", + "aten::upsample_nearest2d", + "aten::upsample_nearest3d" + ), + private=True, +) def _aten_upsample_output_size( self: TReal, output_size: INT64, From 68a14c0917b2acfaac583b6728acd215e712e11e Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Fri, 26 Jan 2024 14:17:10 +0800 Subject: [PATCH 10/13] Update nn.py --- onnxscript/function_libs/torch_lib/ops/nn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index 248e7c26cf..188f8711d6 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2207,7 +2207,7 @@ def _get_upsample_align_corners_mode(align_corners: bool) -> str: "aten::upsample_bilinear2d", "aten::upsample_nearest1d", "aten::upsample_nearest2d", - "aten::upsample_nearest3d" + "aten::upsample_nearest3d", ), private=True, ) From a94559c23d788b03eff14e630867b48d2fd25146 Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Mon, 5 Feb 2024 14:38:29 +0800 Subject: [PATCH 11/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 9 +-- .../function_libs/torch_lib/extra_opinfo.py | 55 +++++++++++++++++++ .../function_libs/torch_lib/ops_test_data.py | 41 +++++++------- 3 files changed, 79 insertions(+), 26 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index 188f8711d6..4e86de7e0f 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2216,7 +2216,6 @@ def _aten_upsample_output_size( output_size: INT64, mode: str, coordinate_transformation_mode: str, - nearest_mode: str = "floor", ) -> TReal: self_shape = op.Shape(self) starts = op.Constant(value_ints=[0]) @@ -2230,7 +2229,6 @@ def _aten_upsample_output_size( output_size, mode=mode, coordinate_transformation_mode=coordinate_transformation_mode, - nearest_mode=nearest_mode, ) @@ -2250,7 +2248,6 @@ def _aten_upsample_scales( None, mode=mode, coordinate_transformation_mode=coordinate_transformation_mode, - nearest_mode="floor", ) @@ -2412,7 +2409,7 @@ def aten_upsample_nearest1d( ) -> TReal: """upsample_nearest1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor""" if size is not None: - return _aten_upsample_output_size(self, size, "nearest", "asymmetric", "floor") + return _aten_upsample_output_size(self, size, "nearest", "asymmetric") else: return _aten_upsample_nearest1d_scales(self, scale_factor) @@ -2460,7 +2457,7 @@ def aten_upsample_nearest2d( del scales_h del scales_w - return _aten_upsample_output_size(self, size, "nearest", "asymmetric", "floor") + return _aten_upsample_output_size(self, size, "nearest", "asymmetric") def aten_upsample_nearest2d_backward( @@ -2489,7 +2486,7 @@ def aten_upsample_nearest3d( del scales_w del scales_d - return _aten_upsample_output_size(self, size, "nearest", "asymmetric", "floor") + return _aten_upsample_output_size(self, size, "nearest", "asymmetric") def aten_upsample_nearest3d_backward( diff --git a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py index 96d38905cd..5e4887ed2f 100644 --- a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py +++ b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py @@ -1626,6 +1626,54 @@ def shape(size, rank, with_batch_channel=True): ) +def sample_inputs_upsample_nearest2d(op_info, device, dtype, requires_grad, **kwargs): + del op_info + del kwargs + + N, C = 2, 3 + D = 4 + SS = 3 + L = 5 + + rank = 2 + + def shape(size, rank, with_batch_channel=True): + if with_batch_channel: + return tuple([N, C] + ([size] * rank)) + return tuple([size] * rank) + + make_arg = functools.partial( + torch_testing.make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + low=-1, + high=1, + ) + + yield opinfo_core.SampleInput(make_arg(shape(D, rank)), shape(SS, rank, False), True) + + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + shape(S, rank, False), + ) + yield opinfo_core.SampleInput( + make_arg(shape(D, rank)), + shape(L, rank, False), + ) + # ONNX don't support below cases: both output_size and scaler are not None + # yield opinfo_core.SampleInput( + # make_arg(shape(D, rank)), + # shape(L, rank, False), + # 1.7, # scaler + # ) + # yield opinfo_core.SampleInput( + # make_arg(shape(D, rank)), + # shape(L, rank, False), + # 0.6, + # ) + + def sample_inputs_upsample_nearest3d(op_info, device, dtype, requires_grad, **kwargs): del op_info del kwargs @@ -2181,6 +2229,13 @@ def __init__(self): sample_inputs_func=sample_inputs_upsample_nearest1d, supports_out=False, ), + opinfo_core.OpInfo( + "ops.aten.upsample_nearest2d", + aten_name="upsample_nearest2d", + dtypes=common_dtype.floating_types_and(torch.bfloat16), + sample_inputs_func=sample_inputs_upsample_nearest2d, + supports_out=False, + ), opinfo_core.OpInfo( "ops.aten.upsample_nearest3d", aten_name="upsample_nearest3d", diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index 46ad8f6a52..15cd5cd3ad 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -2146,20 +2146,21 @@ def _where_input_wrangler( trace_only=True, ), TorchLibOpInfo( - "nn.functional.upsample_nearest2d", + # "nn.functional.upsample_nearest2d", + "ops.aten.upsample_nearest2d", nn_ops.aten_upsample_nearest2d, - input_wrangler=_upsample_input_wrangler, + # input_wrangler=_upsample_input_wrangler, trace_only=True, - ) - .skip( - # Shape should be [N, C, H, W] - matcher=lambda sample: len(sample.input.shape) != 2 + 2, - reason="only test on 2d inputs", - ) - .xfail( - matcher=lambda sample: "scale_factor" in sample.kwargs, - reason="fixme: the scale_factor tests", ), + # .skip( + # # Shape should be [N, C, H, W] + # matcher=lambda sample: len(sample.input.shape) != 2 + 2, + # reason="only test on 2d inputs", + # ) + # .xfail( + # matcher=lambda sample: "scale_factor" in sample.kwargs, + # reason="fixme: the scale_factor tests", + # ), TorchLibOpInfo( "ops.aten.upsample_nearest3d", nn_ops.aten_upsample_nearest3d, @@ -2381,15 +2382,15 @@ def _where_input_wrangler( "nn.functional.celu", ("nn.functional.celu_type_promoted",), ) -ops_test_common.duplicate_opinfo( - OPS_DB, - "nn.functional.upsample_nearest", - ( - # "nn.functional.upsample_nearest1d", - "nn.functional.upsample_nearest2d", - # "nn.functional.upsample_nearest3d", - ), -) +# ops_test_common.duplicate_opinfo( +# OPS_DB, +# "nn.functional.upsample_nearest", +# ( +# # "nn.functional.upsample_nearest1d", +# "nn.functional.upsample_nearest2d", +# # "nn.functional.upsample_nearest3d", +# ), +# ) ops_test_common.duplicate_opinfo( OPS_DB, "ops.aten._log_softmax", ("ops.aten._log_softmax_half",) ) From 11101c9e9bf5de23efe019b10e48f9bb33af51b9 Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Mon, 5 Feb 2024 15:08:28 +0800 Subject: [PATCH 12/13] update --- onnxscript/function_libs/torch_lib/ops/nn.py | 1 + .../function_libs/torch_lib/extra_opinfo.py | 25 ++++++++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index 685c1ce1bc..a6cbb18493 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -2229,6 +2229,7 @@ def _aten_upsample_output_size( output_size, mode=mode, coordinate_transformation_mode=coordinate_transformation_mode, + nearest_mode="floor", ) diff --git a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py index 04cb055b37..086264e9bf 100644 --- a/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py +++ b/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py @@ -1683,7 +1683,6 @@ def sample_inputs_upsample_nearest3d(op_info, device, dtype, requires_grad, **kw SS = 3 L = 5 - align_corners_options = (True, False) rank = 3 def shape(size, rank, with_batch_channel=True): @@ -1724,6 +1723,30 @@ def shape(size, rank, with_batch_channel=True): def sample_inputs_upsample_trilinear3d(op_info, device, dtype, requires_grad, **kwargs): + del op_info + del kwargs + + N, C = 2, 3 + D = 4 + SS = 3 + L = 5 + + align_corners_options = (True, False) + rank = 3 + + def shape(size, rank, with_batch_channel=True): + if with_batch_channel: + return tuple([N, C] + ([size] * rank)) + return tuple([size] * rank) + + make_arg = functools.partial( + torch_testing.make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + low=-1, + high=1, + ) for align_corners in align_corners_options: yield opinfo_core.SampleInput( From 221d313b4d7ef40d82be05177ab6d26d4f02051a Mon Sep 17 00:00:00 2001 From: xiaowuhu Date: Mon, 5 Feb 2024 15:27:35 +0800 Subject: [PATCH 13/13] Update ops_test_data.py --- .../function_libs/torch_lib/ops_test_data.py | 32 ------------------- 1 file changed, 32 deletions(-) diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py index b63f660e04..3cfd4a1629 100644 --- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py +++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py @@ -415,18 +415,6 @@ def _sum_input_wrangler( return args, kwargs -def _upsample_input_wrangler( - args: list[Any], kwargs: dict[str, Any] -) -> tuple[list[Any], dict[str, Any]]: - if "scale_factor" in kwargs: - kwargs["scales_h"] = kwargs["scale_factor"] - kwargs["scales_w"] = kwargs["scale_factor"] - del kwargs["scale_factor"] - if "size" in kwargs: - kwargs["size"] = np.array(kwargs["size"], dtype=np.int64) - return args, kwargs - - def _unflatten_input_wrangler( args: list[Any], kwargs: dict[str, Any] ) -> tuple[list[Any], dict[str, Any]]: @@ -2146,21 +2134,10 @@ def _where_input_wrangler( trace_only=True, ), TorchLibOpInfo( - # "nn.functional.upsample_nearest2d", "ops.aten.upsample_nearest2d", nn_ops.aten_upsample_nearest2d, - # input_wrangler=_upsample_input_wrangler, trace_only=True, ), - # .skip( - # # Shape should be [N, C, H, W] - # matcher=lambda sample: len(sample.input.shape) != 2 + 2, - # reason="only test on 2d inputs", - # ) - # .xfail( - # matcher=lambda sample: "scale_factor" in sample.kwargs, - # reason="fixme: the scale_factor tests", - # ), TorchLibOpInfo( "ops.aten.upsample_nearest3d", nn_ops.aten_upsample_nearest3d, @@ -2387,15 +2364,6 @@ def _where_input_wrangler( "nn.functional.celu", ("nn.functional.celu_type_promoted",), ) -# ops_test_common.duplicate_opinfo( -# OPS_DB, -# "nn.functional.upsample_nearest", -# ( -# # "nn.functional.upsample_nearest1d", -# "nn.functional.upsample_nearest2d", -# # "nn.functional.upsample_nearest3d", -# ), -# ) ops_test_common.duplicate_opinfo( OPS_DB, "ops.aten._log_softmax", ("ops.aten._log_softmax_half",) )