Merge branch 'main' into xiaowu/AddOp(upsample_linear1d)

Author: justinchuby

noxfile.py

@@ -29,6 +29,7 @@
 ONNX = "onnx==1.14.1"
 ONNX_RUNTIME = "onnxruntime==1.16.1"
 PYTORCH = "torch==2.1.0"
+TORCHVISON = "torchvision==0.16"
 ONNX_RUNTIME_NIGHTLY_DEPENDENCIES = (
     "flatbuffers",
     "coloredlogs",
@@ -52,6 +53,7 @@ def test(session):
     session.install(
         *COMMON_TEST_DEPENDENCIES,
         PYTORCH,
+        TORCHVISON,
         ONNX,
         ONNX_RUNTIME,
     )
@@ -78,7 +80,7 @@ def test_torch_nightly(session):
 @nox.session(tags=["test-onnx-weekly"])
 def test_onnx_weekly(session):
     """Test with ONNX weekly (preview) build."""
-    session.install(*COMMON_TEST_DEPENDENCIES, ONNX_RUNTIME, PYTORCH)
+    session.install(*COMMON_TEST_DEPENDENCIES, ONNX_RUNTIME, PYTORCH, TORCHVISON)
     session.install("-r", "requirements/ci/requirements-onnx-weekly.txt")
     session.install(".", "--no-deps")
     session.run("pip", "list")
@@ -89,7 +91,11 @@ def test_onnx_weekly(session):
 def test_ort_nightly(session):
     """Test with ONNX Runtime nightly builds."""
     session.install(
-        *COMMON_TEST_DEPENDENCIES, PYTORCH, ONNX, *ONNX_RUNTIME_NIGHTLY_DEPENDENCIES
+        *COMMON_TEST_DEPENDENCIES,
+        PYTORCH,
+        TORCHVISON,
+        ONNX,
+        *ONNX_RUNTIME_NIGHTLY_DEPENDENCIES,
     )
     session.install("-r", "requirements/ci/requirements-ort-nightly.txt")
     session.install(".", "--no-deps")
@@ -101,7 +107,11 @@ def test_ort_nightly(session):
 def test_experimental_torchlib_tracing(session):
     """Test TorchLib with the experimental TORCHLIB_EXPERIMENTAL_PREFER_TRACING flag on."""
     session.install(
-        *COMMON_TEST_DEPENDENCIES, PYTORCH, ONNX, *ONNX_RUNTIME_NIGHTLY_DEPENDENCIES
+        *COMMON_TEST_DEPENDENCIES,
+        PYTORCH,
+        TORCHVISON,
+        ONNX,
+        *ONNX_RUNTIME_NIGHTLY_DEPENDENCIES,
     )
     session.install("-r", "requirements/ci/requirements-ort-nightly.txt")
     session.install(".", "--no-deps")

onnxscript/function_libs/torch_lib/ops/__init__.py

@@ -7,6 +7,7 @@
     "prims",
     "sparse",
     "special",
+    "vision",
 ]
 
-from . import core, fft, linalg, nested, nn, prims, sparse, special
+from . import core, fft, linalg, nested, nn, prims, sparse, special, vision

onnxscript/function_libs/torch_lib/ops/nn.py

@@ -2197,85 +2197,99 @@ def aten_unflatten_dense_tensors(
     raise NotImplementedError()
 
 
-@torch_op(("aten::upsample_bicubic2d", "aten::upsample_bicubic2d.vec"), trace_only=True)
-def aten_upsample_bicubic2d(
-    self: TReal,
-    output_size: INT64,
-    align_corners: bool,
-    scale_factors: Optional[TFloat] = None,
-) -> TReal:
-    """upsample_bicubic2d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor
-    upsample_bicubic2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor
-    """
-
-    if output_size is not None:
-        result = _aten_upsample_output_size(self, output_size, align_corners, "cubic")
-    else:
-        result = _aten_upsample_scales(self, scale_factors, align_corners, "cubic")
-    return result
+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", private=True)
+@torch_op(("aten::upsample_bicubic2d", "aten::upsample_bilinear2d"), private=True)
 def _aten_upsample_output_size(
     self: TReal,
     output_size: INT64,
-    align_corners: bool,
-    str_mode: str,
+    mode: str,
+    coordinate_transformation_mode: str,
 ) -> 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, output_size, axis=0)
-    if align_corners:
-        result = op.Resize(
-            self,
-            None,
-            None,
-            output_size,
-            mode=str_mode,
-            coordinate_transformation_mode="align_corners",
-        )
-    else:
-        result = op.Resize(
-            self,
-            None,
-            None,
-            output_size,
-            mode=str_mode,
-            coordinate_transformation_mode="pytorch_half_pixel",
-        )
-
-    return result
+    return op.Resize(
+        self,
+        None,
+        None,
+        output_size,
+        mode=mode,
+        coordinate_transformation_mode=coordinate_transformation_mode,
+        nearest_mode="floor",
+    )
 
 
-@torch_op("aten::upsample_bicubic2d", private=True)
+@torch_op(("aten::upsample_bicubic2d", "aten::upsample_bilinear2d"), private=True)
 def _aten_upsample_scales(
     self: TReal,
     scale_factors: TFloat,
-    align_corners: bool,
-    str_mode: str,
+    mode: str,
+    coordinate_transformation_mode: str,
 ) -> 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)
-    if align_corners:
-        result = op.Resize(
+    return op.Resize(
+        self,
+        None,
+        scale_factors,  # format should be: [1.0, 1.0, scale_h, scale_w]
+        None,
+        mode=mode,
+        coordinate_transformation_mode=coordinate_transformation_mode,
+        nearest_mode="floor",
+    )
+
+
+@torch_op("aten::upsample_bicubic2d", trace_only=True)
+def aten_upsample_bicubic2d(
+    self: TReal,
+    output_size: INT64,
+    align_corners: bool,
+    scales_h: Optional[float] = None,
+    scales_w: Optional[float] = None,
+) -> TReal:
+    """upsample_bicubic2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor"""
+
+    # NOTE: Based on experimentation, scales_h and scales_w are always ignored in PyTorch,
+    # unless when align_corners is True, in which case we do not know what is going on.
+    coordinate_transformation_mode = _get_upsample_align_corners_mode(align_corners)
+    return _aten_upsample_output_size(
+        self,
+        output_size,
+        mode="cubic",
+        coordinate_transformation_mode=coordinate_transformation_mode,
+    )
+
+
+@torch_op("aten::upsample_bicubic2d.vec", trace_only=True)
+def aten_upsample_bicubic2d_vec(
+    self: TReal,
+    output_size: INT64,
+    align_corners: bool,
+    scale_factors: Optional[Sequence[float]],
+) -> TReal:
+    """upsample_bicubic2d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor"""
+
+    coordinate_transformation_mode = _get_upsample_align_corners_mode(align_corners)
+    if scale_factors is not None:
+        result = _aten_upsample_scales(
             self,
-            None,
-            scale_factors,  # format should be: [1.0, 1.0, scale_h, scale_w]
-            None,
-            mode=str_mode,
-            coordinate_transformation_mode="align_corners",
+            op.Constant(value_floats=scale_factors),
+            mode="cubic",
+            coordinate_transformation_mode=coordinate_transformation_mode,
         )
     else:
-        result = op.Resize(
+        result = _aten_upsample_output_size(
             self,
-            None,
-            scale_factors,  # format should be: [1.0, 1.0, scale_h, scale_w]
-            None,
-            mode=str_mode,
-            coordinate_transformation_mode="pytorch_half_pixel",
+            output_size,
+            mode="cubic",
+            coordinate_transformation_mode=coordinate_transformation_mode,
         )
+
     return result
 
 
@@ -2295,67 +2309,50 @@ def aten_upsample_bicubic2d_backward(
 @torch_op("aten::upsample_bilinear2d", trace_only=True)
 def aten_upsample_bilinear2d(
     self: TReal,
-    output_size: Optional[INT64] = None,
+    output_size: INT64,
+    align_corners: bool,
     scales_h: Optional[float] = None,
     scales_w: Optional[float] = None,
-    align_corners: bool = True,  # pylint: disable=unused-argument
-) -> TReal:
-    """upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor"""
-
-    if output_size is not None:
-        result = _aten_upsample_bilinear2d_output_size(self, output_size)
-    else:
-        assert scales_h is not None
-        assert scales_h == scales_w
-        result = _aten_upsample_bilinear2d_scales(self, scales_h, scales_w)
-    return result
-
-
-@torch_op("aten::upsample_bilinear2d", private=True)
-def _aten_upsample_bilinear2d_output_size(
-    self: TReal,
-    output_size: INT64,
 ) -> TReal:
     """upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor"""
 
-    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, output_size, axis=0)
-    return op.Resize(
+    # NOTE: Based on experimentation, scales_h and scales_w are always ignored in PyTorch,
+    # unless when align_corners is True, in which case we do not know what is going on.
+    coordinate_transformation_mode = _get_upsample_align_corners_mode(align_corners)
+    return _aten_upsample_output_size(
         self,
-        None,
-        None,
         output_size,
+        coordinate_transformation_mode=coordinate_transformation_mode,
         mode="linear",
-        coordinate_transformation_mode="align_corners",
     )
 
 
-@torch_op("aten::upsample_bilinear2d", private=True)
-def _aten_upsample_bilinear2d_scales(
+@torch_op("aten::upsample_bilinear2d.vec", trace_only=True)
+def aten_upsample_bilinear2d_vec(
     self: TReal,
-    scales_h: float,
-    scales_w: float,
+    output_size: Optional[INT64],
+    align_corners: bool,
+    scale_factors: Optional[Sequence[float]],
 ) -> TReal:
-    """upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor"""
+    """upsample_bilinear2d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor"""
 
-    neg_1 = op.Constant(value_ints=[-1])
-    scales = op.Concat(
-        op.Constant(value_floats=[1.0, 1.0]),
-        op.Reshape(op.Constant(value_float=scales_h), neg_1),
-        op.Reshape(op.Constant(value_float=scales_w), neg_1),
-        axis=0,
-    )
-    return op.Resize(
-        self,
-        None,
-        scales,  # format should be: [1.0, 1.0, scale_h, scale_w]
-        None,
-        mode="linear",
-        coordinate_transformation_mode="align_corners",
-    )
+    coordinate_transformation_mode = _get_upsample_align_corners_mode(align_corners)
+    if scale_factors is not None:
+        result = _aten_upsample_scales(
+            self,
+            op.Constant(value_floats=scale_factors),
+            mode="linear",
+            coordinate_transformation_mode=coordinate_transformation_mode,
+        )
+    else:
+        result = _aten_upsample_output_size(
+            self,
+            output_size,
+            mode="linear",
+            coordinate_transformation_mode=coordinate_transformation_mode,
+        )
+
+    return result
 
 
 def aten_upsample_bilinear2d_backward(

onnxscript/function_libs/torch_lib/ops/vision.py

@@ -0,0 +1,24 @@
+# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value"
+"""torchvision operators."""
+from __future__ import annotations
+
+from onnxscript.function_libs.torch_lib.registration import torch_op
+from onnxscript.onnx_opset import opset18 as op
+from onnxscript.onnx_types import FLOAT, INT64
+
+_INT64_MAX = 0x7FFFFFFFFFFFFFFF
+
+
+@torch_op("torchvision::nms")
+def torchvision_nms(boxes: FLOAT, scores: FLOAT, iou_threshold: float) -> INT64:
+    # boxes: [num_batches, spatial_dimension, 4]
+    boxes = op.Unsqueeze(boxes, [0])
+    # scores: [num_batches, num_classes, spatial_dimension]
+    scores = op.Unsqueeze(scores, [0, 1])
+    # nms_out: [num_selected_indices, 3] where each column is [batch_index, class_index, box_index]
+    nms_out = op.NonMaxSuppression(boxes, scores, _INT64_MAX, iou_threshold)
+    return op.Reshape(op.Slice(nms_out, axes=[1], starts=[2], ends=[3]), [-1])