Inconsistency between RoIPool's forward and roi_pool #6395
Description
🚀 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 "", 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