Skip to content

Commit

Permalink
Make 'ROIAlign' & 'ROIAlignV2' version of ROIPooler scriptable.
Browse files Browse the repository at this point in the history
Summary: Pull Request resolved: #1835

Reviewed By: rbgirshick

Differential Revision: D22819550

Pulled By: ppwwyyxx

fbshipit-source-id: 85cd2198676289e0ab02678f221b97887e543395
  • Loading branch information
chenbohua3 authored and facebook-github-bot committed Jul 30, 2020
1 parent af866c4 commit be792b9
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 15 deletions.
42 changes: 27 additions & 15 deletions detectron2/modeling/poolers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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).
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
44 changes: 44 additions & 0 deletions tests/modeling/test_roi_pooler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -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()

0 comments on commit be792b9

Please sign in to comment.