Refactor sample_inputs_roi_align for improved clarity and efficiency; remove redundant tests and simplify input handling

Author: FraGirla

tests/function_libs/torch_lib/extra_opinfo.py

@@ -1471,81 +1471,34 @@ def sample_inputs_replication_pad1d(op_info, device, dtype, requires_grad, **kwa
 
 
 def sample_inputs_roi_align(op_info, device, dtype, requires_grad, **kwargs):
-    del op_info
-    del kwargs
-    # roi_align signature: (input, boxes, output_size, spatial_scale=1.0, sampling_ratio=-1, aligned=False)
-
-    # Test 1: spatial_scale=1, sampling_ratio=2
-    x1 = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi1 = torch.tensor([[0, 1.5, 1.5, 3, 3]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x1,
-        args=(roi1, (5, 5)),
-        kwargs={"spatial_scale": 1.0, "sampling_ratio": 2, "aligned": True},
-    )
-
-    # Test 2: spatial_scale=0.5, sampling_ratio=3
-    x2 = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi2 = torch.tensor([[0, 0.2, 0.3, 4.5, 3.5]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x2,
-        args=(roi2, (5, 5)),
-        kwargs={"spatial_scale": 0.5, "sampling_ratio": 3, "aligned": True},
-    )
-
-    # Test 3: spatial_scale=1.8, sampling_ratio=2
-    x3 = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi3 = torch.tensor([[0, 0.2, 0.3, 4.5, 3.5]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x3,
-        args=(roi3, (5, 5)),
-        kwargs={"spatial_scale": 1.8, "sampling_ratio": 2, "aligned": True},
-    )
-
-    # Test 4: spatial_scale=2.5, sampling_ratio=0, output_size=(2,2)
-    x4 = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi4 = torch.tensor([[0, 0.2, 0.3, 4.5, 3.5]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x4,
-        args=(roi4, (2, 2)),
-        kwargs={"spatial_scale": 2.5, "sampling_ratio": 0, "aligned": True},
-    )
+    del op_info, kwargs
 
-    # Test 5: spatial_scale=2.5, sampling_ratio=-1, output_size=(2,2)
-    x5 = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi5 = torch.tensor([[0, 0.2, 0.3, 4.5, 3.5]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x5,
-        args=(roi5, (2, 2)),
-        kwargs={"spatial_scale": 2.5, "sampling_ratio": -1, "aligned": True},
-    )
+    def make_x():
+        return torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
 
-    # Test 6: malformed boxes (test_roi_align_malformed_boxes)
-    x6 = torch.randn(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi6 = torch.tensor([[0, 2, 0.3, 1.5, 1.5]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x6,
-        args=(roi6, (5, 5)),
-        kwargs={"spatial_scale": 1.0, "sampling_ratio": 1, "aligned": True},
-    )
+    # rois is [K, 5] = [batch_idx, x1, y1, x2, y2]
+    roi_a = torch.tensor([[0, 1.5, 1.5, 3.0, 3.0]], dtype=dtype, device=device)
+    roi_b = torch.tensor([[0, 0.2, 0.3, 4.5, 3.5]], dtype=dtype, device=device)
+    roi_int = torch.tensor([[0, 0.0, 0.0, 4.0, 4.0]], dtype=dtype, device=device)
+    roi_malformed = torch.tensor([[0, 2.0, 0.3, 1.5, 1.5]], dtype=dtype, device=device)  # x1 > x2-ish
 
-    # Test 7: aligned=False, spatial_scale=1, sampling_ratio=2
-    x7 = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi7 = torch.tensor([[0, 0, 0, 4, 4]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x7,
-        args=(roi7, (5, 5)),
-        kwargs={"spatial_scale": 1.0, "sampling_ratio": 2, "aligned": False},
-    )
+    # (rois, spatial_scale, pooled_h, pooled_w, sampling_ratio, aligned)
+    cases = [
+        (roi_a, 1.0, 5, 5, 2, True),
+        (roi_b, 0.5, 5, 5, 3, True),
+        (roi_b, 1.8, 5, 5, 2, True),
+        (roi_b, 2.5, 2, 2, 0, True),
+        (roi_b, 2.5, 2, 2, -1, True),
+        (roi_malformed, 1.0, 5, 5, 1, True),
+        (roi_int, 1.0, 5, 5, 2, False),
+        (roi_int, 1.0, 5, 5, -1, False),
+    ]
 
-    # Test 8: aligned=False, spatial_scale=1, sampling_ratio=-1
-    x8 = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=requires_grad)
-    roi8 = torch.tensor([[0, 0, 0, 4, 4]], dtype=dtype, device=device)
-    yield opinfo_core.SampleInput(
-        x8,
-        args=(roi8, (5, 5)),
-        kwargs={"spatial_scale": 1.0, "sampling_ratio": -1, "aligned": False},
-    )
+    for rois, spatial_scale, ph, pw, sr, aligned in cases:
+        yield opinfo_core.SampleInput(
+            make_x(),
+            args=(rois, float(spatial_scale), int(ph), int(pw), int(sr), bool(aligned)),
+        )
 
 
 def sample_inputs_roi_pool(op_info, device, dtype, requires_grad, **kwargs):
@@ -3132,7 +3085,7 @@ def __init__(self):
     ),
     opinfo_core.OpInfo(
         "torchvision.ops.roi_align",
-        op=torchvision.ops.roi_align,
+        op=torch.ops.torchvision.roi_align.default,
         dtypes=common_dtype.floating_types(),
         sample_inputs_func=sample_inputs_roi_align,
         supports_out=False,

tests/function_libs/torch_lib/ops_test_data.py

@@ -448,36 +448,6 @@ def _where_input_wrangler(
     args[0], args[1] = args[1], args[0]
     return args, kwargs
 
-
-def _torchvision_roi_align_default_input_wrangler(
-    args: list[Any], kwargs: dict[str, Any]
-) -> tuple[list[Any], dict[str, Any]]:
-    # Convert:
-    #   roi_align(input, boxes, output_size, spatial_scale=..., sampling_ratio=..., aligned=...)
-    # into:
-    #   roi_align(input, boxes, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned)
-    output_size = args.pop(2)
-    if isinstance(output_size, np.ndarray):
-        if output_size.ndim == 0:
-            pooled_height = int(output_size)
-            pooled_width = int(output_size)
-        else:
-            pooled_height, pooled_width = output_size.tolist()
-    elif isinstance(output_size, (tuple, list)):
-        pooled_height, pooled_width = output_size
-    else:
-        pooled_height = output_size
-        pooled_width = output_size
-
-    pooled_height = int(pooled_height)
-    pooled_width = int(pooled_width)
-    spatial_scale = float(kwargs.pop("spatial_scale", 1.0))
-    sampling_ratio = int(kwargs.pop("sampling_ratio", -1))
-    aligned = bool(kwargs.pop("aligned", False))
-    args.extend([spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned])
-    return args, {}
-
-
 # Ops to be tested for numerical consistency between onnx and pytorch
 # Find the names of the OpInfos in torch/testing/_internal/common_methods_invocations.py
 TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
@@ -1948,11 +1918,7 @@ def _torchvision_roi_align_default_input_wrangler(
     ),
     TorchLibOpInfo("zeros_like", core_ops.aten_zeros_like),
     TorchLibOpInfo("torchvision.ops.nms", vision_ops.torchvision_nms),
-    TorchLibOpInfo(
-        "torchvision.ops.roi_align",
-        vision_ops.torchvision_roi_align,
-        input_wrangler=_torchvision_roi_align_default_input_wrangler,
-    ),
+    TorchLibOpInfo("torchvision.ops.roi_align", vision_ops.torchvision_roi_align),
     TorchLibOpInfo("torchvision.ops.roi_pool", vision_ops.torchvision_roi_pool),
 )