Repository: pytorch/vision

Issue:
Inconsistency between RoIPool's forward and roi_pool
### 🚀 The feature

Currently the functional `roi_pool` method supports both `Tensor` and `List[Tensor]` for boxes, but the module forward function of `RoIPool` expects boxes(rois) to be `Tensor` only. Can we change the function signature to support both? 

The inconsistency can cause issue during torchscripting, for example:
```
from typing import List
import torch
from torch import Tensor
from torchvision.ops import RoIPool

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.pool = RoIPool(output_size=[3,3], spatial_scale=1.0)

    def forward(self, imgs, boxes: List[Tensor]):
        return self.pool(imgs, boxes)

model = Model()
torch.jit.script(model)
```
Raise the following error:

> forward(__torch__.torchvision.ops.roi_pool.RoIPool self, Tensor input, Tensor rois) -> Tensor:
> Expected a value of type 'Tensor' for argument 'rois' but instead found type 'List[Tensor]'.
> Empty lists default to List[Tensor]. Add a variable annotation to the assignment to create an empty list of another type (torch.jit.annotate(List[T, []]) where T is the type of elements in the list for Python 2)
> :
>   File "<ipython-input-8-9536dd2f0cc1>", line 12
>     def forward(self, imgs, boxes: List[Tensor]):
>         return self.pool(imgs, boxes)
>                ~~~~~~~~~ <--- HERE

### Motivation, pitch

Make sure the model can be scriptable when using list of boxes as the input.

### Alternatives

Only support Tensor as the input, user are required to convert bbox of different images into one tensor. Potentially make `_utils.convert_boxes_to_roi_format` public so that users don't need to write the conversion function.

### Additional context

_No response_


Discussion:

Resolving patch:
diff --git a/torchvision/ops/roi_align.py b/torchvision/ops/roi_align.py
--- a/torchvision/ops/roi_align.py
+++ b/torchvision/ops/roi_align.py
@@ -82,7 +82,7 @@ def __init__(
         self.sampling_ratio = sampling_ratio
         self.aligned = aligned
 
-    def forward(self, input: Tensor, rois: Tensor) -> Tensor:
+    def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor:
         return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, self.aligned)
 
     def __repr__(self) -> str:
diff --git a/torchvision/ops/roi_pool.py b/torchvision/ops/roi_pool.py
--- a/torchvision/ops/roi_pool.py
+++ b/torchvision/ops/roi_pool.py
@@ -62,7 +62,7 @@ def __init__(self, output_size: BroadcastingList2[int], spatial_scale: float):
         self.output_size = output_size
         self.spatial_scale = spatial_scale
 
-    def forward(self, input: Tensor, rois: Tensor) -> Tensor:
+    def forward(self, input: Tensor, rois: Union[Tensor, List[Tensor]]) -> Tensor:
         return roi_pool(input, rois, self.output_size, self.spatial_scale)
 
     def __repr__(self) -> str:


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -69,6 +69,15 @@ def forward(self, a):
         self.layer(a)
 
 
+class PoolWrapper(nn.Module):
+    def __init__(self, pool: nn.Module):
+        super().__init__()
+        self.pool = pool
+
+    def forward(self, imgs: Tensor, boxes: List[Tensor]) -> Tensor:
+        return self.pool(imgs, boxes)
+
+
 class RoIOpTester(ABC):
     dtype = torch.float64
 
@@ -150,6 +159,14 @@ def _helper_boxes_shape(self, func):
             boxes = torch.tensor([[0, 0, 3]], dtype=a.dtype)
             ops.roi_pool(a, [boxes], output_size=(2, 2))
 
+    def _helper_jit_boxes_list(self, model):
+        x = torch.rand(2, 1, 10, 10)
+        roi = torch.tensor([[0, 0, 0, 9, 9], [0, 0, 5, 4, 9], [0, 5, 5, 9, 9], [1, 0, 0, 9, 9]], dtype=torch.float).t()
+        rois = [roi, roi]
+        scriped = torch.jit.script(model)
+        y = scriped(x, rois)
+        assert y.shape == (10, 1, 3, 3)
+
     @abstractmethod
     def fn(*args, **kwargs):
         pass
@@ -210,6 +227,10 @@ def get_slice(k, block):
     def test_boxes_shape(self):
         self._helper_boxes_shape(ops.roi_pool)
 
+    def test_jit_boxes_list(self):
+        model = PoolWrapper(ops.RoIPool(output_size=[3, 3], spatial_scale=1.0))
+        self._helper_jit_boxes_list(model)
+
 
 class TestPSRoIPool(RoIOpTester):
     def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs):
@@ -450,6 +471,10 @@ def test_qroi_align_multiple_images(self):
         with pytest.raises(RuntimeError, match="Only one image per batch is allowed"):
             ops.roi_align(qx, qrois, output_size=5)
 
+    def test_jit_boxes_list(self):
+        model = PoolWrapper(ops.RoIAlign(output_size=[3, 3], spatial_scale=1.0, sampling_ratio=-1))
+        self._helper_jit_boxes_list(model)
+
 
 class TestPSRoIAlign(RoIOpTester):
     def fn(self, x, rois, pool_h, pool_w, spatial_scale=1, sampling_ratio=-1, **kwargs):