Make 'ROIAlign' & 'ROIAlignV2' version of ROIPooler scriptable.

Summary: Pull Request resolved: #1835

Reviewed By: rbgirshick

Differential Revision: D22819550

Pulled By: ppwwyyxx

fbshipit-source-id: 85cd2198676289e0ab02678f221b97887e543395

detectron2/modeling/poolers.py

@@ -7,12 +7,25 @@
 from torchvision.ops import RoIPool
 
 from detectron2.layers import ROIAlign, ROIAlignRotated, cat, nonzero_tuple
+from detectron2.structures import Boxes
+
+"""
+To export ROIPooler to torchscript, in this file, variables that should be annotated with
+`Union[List[Boxes], List[RotatedBoxes]]` are only annotated with `List[Boxes]`.
+
+TODO: Correct these annotations when torchscript support `Union`.
+https://github.com/pytorch/pytorch/issues/41412
+"""
 
 __all__ = ["ROIPooler"]
 
 
 def assign_boxes_to_levels(
- box_lists, min_level: int, max_level: int, canonical_box_size: int, canonical_level: int
+ box_lists: List[Boxes],
+ min_level: int,
+ max_level: int,
+ canonical_box_size: int,
+ canonical_level: int,
 ):
  """
  Map each box in `box_lists` to a feature map level index and return the assignment
@@ -35,19 +48,25 @@ def assign_boxes_to_levels(
  `self.min_level`, for the corresponding box (so value i means the box is at
  `self.min_level + i`).
  """
- eps = sys.float_info.epsilon
  box_sizes = torch.sqrt(cat([boxes.area() for boxes in box_lists]))
  # Eqn.(1) in FPN paper
  level_assignments = torch.floor(
- canonical_level + torch.log2(box_sizes / canonical_box_size + eps)
+ canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8)
  )
  # clamp level to (min, max), in case the box size is too large or too small
  # for the available feature maps
  level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level)
  return level_assignments.to(torch.int64) - min_level
 
 
-def convert_boxes_to_pooler_format(box_lists):
+def _fmt_box_list(box_tensor, batch_index: int):
+ repeated_index = torch.full(
+ (len(box_tensor), 1), batch_index, dtype=box_tensor.dtype, device=box_tensor.device
+ )
+ return cat((repeated_index, box_tensor), dim=1)
+
+
+def convert_boxes_to_pooler_format(box_lists: List[Boxes]):
  """
  Convert all boxes in `box_lists` to the low-level format used by ROI pooling ops
  (see description under Returns).
@@ -70,15 +89,8 @@ def convert_boxes_to_pooler_format(box_lists):
  where batch index is the index in [0, N) identifying which batch image the
  rotated box (x_ctr, y_ctr, width, height, angle_degrees) comes from.
  """
-
- def fmt_box_list(box_tensor, batch_index):
- repeated_index = torch.full(
- (len(box_tensor), 1), batch_index, dtype=box_tensor.dtype, device=box_tensor.device
- )
- return cat((repeated_index, box_tensor), dim=1)
-
  pooler_fmt_boxes = cat(
- [fmt_box_list(box_list.tensor, i) for i, box_list in enumerate(box_lists)], dim=0
+ [_fmt_box_list(box_list.tensor, i) for i, box_list in enumerate(box_lists)], dim=0
  )
 
  return pooler_fmt_boxes
@@ -176,7 +188,7 @@ def __init__(
  assert canonical_box_size > 0
  self.canonical_box_size = canonical_box_size
 
- def forward(self, x: List[torch.Tensor], box_lists):
+ def forward(self, x: List[torch.Tensor], box_lists: List[Boxes]):
  """
  Args:
  x (list[Tensor]): A list of feature maps of NCHW shape, with scales matching those
@@ -226,9 +238,9 @@ def forward(self, x: List[torch.Tensor], box_lists):
  (num_boxes, num_channels, output_size, output_size), dtype=dtype, device=device
  )
 
- for level, (x_level, pooler) in enumerate(zip(x, self.level_poolers)):
+ for level, pooler in enumerate(self.level_poolers):
  inds = nonzero_tuple(level_assignments == level)[0]
  pooler_fmt_boxes_level = pooler_fmt_boxes[inds]
- output[inds] = pooler(x_level, pooler_fmt_boxes_level)
+ output[inds] = pooler(x[level], pooler_fmt_boxes_level)
 
  return output

tests/modeling/test_roi_pooler.py

@@ -5,6 +5,7 @@
 
 from detectron2.modeling.poolers import ROIPooler
 from detectron2.structures import Boxes, RotatedBoxes
+from detectron2.utils.env import TORCH_VERSION
 
 logger = logging.getLogger(__name__)
 
@@ -80,6 +81,49 @@ def test_roialignv2_roialignrotated_match_cpu(self):
  def test_roialignv2_roialignrotated_match_cuda(self):
  self._test_roialignv2_roialignrotated_match(device="cuda")
 
+ def _test_scriptability(self, device):
+ pooler_resolution = 14
+ canonical_level = 4
+ canonical_scale_factor = 2 ** canonical_level
+ pooler_scales = (1.0 / canonical_scale_factor,)
+ sampling_ratio = 0
+
+ N, C, H, W = 2, 4, 10, 8
+ N_rois = 10
+ std = 11
+ mean = 0
+ feature = (torch.rand(N, C, H, W) - 0.5) * 2 * std + mean
+
+ features = [feature.to(device)]
+
+ rois = []
+ for _ in range(N):
+ boxes = self._rand_boxes(
+ num_boxes=N_rois, x_max=W * canonical_scale_factor, y_max=H * canonical_scale_factor
+ )
+
+ rois.append(Boxes(boxes).to(device))
+
+ roialignv2_pooler = ROIPooler(
+ output_size=pooler_resolution,
+ scales=pooler_scales,
+ sampling_ratio=sampling_ratio,
+ pooler_type="ROIAlignV2",
+ )
+
+ roialignv2_out = roialignv2_pooler(features, rois)
+ scripted_roialignv2_out = torch.jit.script(roialignv2_pooler)(features, rois)
+ self.assertTrue(torch.equal(roialignv2_out, scripted_roialignv2_out))
+
+ @unittest.skipIf(TORCH_VERSION < (1, 7), "Insufficient pytorch version")
+ def test_scriptability_cpu(self):
+ self._test_scriptability(device="cpu")
+
+ @unittest.skipIf(TORCH_VERSION < (1, 7), "Insufficient pytorch version")
+ @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
+ def test_scriptability_gpu(self):
+ self._test_scriptability(device="cuda")
+
 
 if __name__ == "__main__":
  unittest.main()