add shapes_to_tensor that stacks scalars or ints in a traceable/scriptable way

Summary: The function is reused in 3 different places.

Reviewed By: zhanghang1989

Differential Revision: D30527464

fbshipit-source-id: 90b0fdbabd296ba5f4d911a8481c7c454b947b7b

Author: ppwwyyxx

detectron2/layers/__init__.py

@@ -15,6 +15,7 @@
     Linear,
     nonzero_tuple,
     cross_entropy,
+    shapes_to_tensor,
 )
 from .blocks import CNNBlockBase, DepthwiseSeparableConv2d
 from .aspp import ASPP

detectron2/layers/wrappers.py

@@ -8,11 +8,33 @@
 is implemented
 """
 
-from typing import List
+from typing import List, Optional
 import torch
 from torch.nn import functional as F
 
 
+def shapes_to_tensor(x: List[int], device: Optional[torch.device] = None) -> torch.Tensor:
+    """
+    Turn a list of integer scalars or integer Tensor scalars into a vector,
+    in a way that's both traceable and scriptable.
+
+    In tracing, `x` should be a list of scalar Tensor, so the output can trace to the inputs.
+    In scripting or eager, `x` should be a list of int.
+    """
+    if torch.jit.is_scripting():
+        return torch.as_tensor(x, device=device)
+    if torch.jit.is_tracing():
+        assert all(
+            [isinstance(t, torch.Tensor) for t in x]
+        ), "Shape should be tensor during tracing!"
+        # as_tensor should not be used in tracing because it records a constant
+        ret = torch.stack(x)
+        if ret.device != device:  # avoid recording a hard-coded device if not necessary
+            ret = ret.to(device=device)
+        return ret
+    return torch.as_tensor(x, device=device)
+
+
 def cat(tensors: List[torch.Tensor], dim: int = 0):
     """
     Efficient version of torch.cat that avoids a copy if there is only a single element in a list

detectron2/modeling/poolers.py

@@ -5,7 +5,7 @@
 from torch import nn
 from torchvision.ops import RoIPool
 
-from detectron2.layers import ROIAlign, ROIAlignRotated, cat, nonzero_tuple
+from detectron2.layers import ROIAlign, ROIAlignRotated, cat, nonzero_tuple, shapes_to_tensor
 from detectron2.structures import Boxes
 
 """
@@ -58,13 +58,6 @@ def assign_boxes_to_levels(
     return level_assignments.to(torch.int64) - min_level
 
 
-def _fmt_box_list(box_tensor, batch_index: int):
-    repeated_index = torch.full_like(
-        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
@@ -88,11 +81,13 @@ def convert_boxes_to_pooler_format(box_lists: List[Boxes]):
             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.
     """
-    pooler_fmt_boxes = cat(
-        [_fmt_box_list(box_list.tensor, i) for i, box_list in enumerate(box_lists)], dim=0
+    boxes = torch.cat([x.tensor for x in box_lists], dim=0)
+    # __len__ returns Tensor in tracing.
+    sizes = shapes_to_tensor([x.__len__() for x in box_lists], device=boxes.device)
+    indices = torch.repeat_interleave(
+        torch.arange(len(box_lists), dtype=boxes.dtype, device=boxes.device), sizes
     )
-
-    return pooler_fmt_boxes
+    return cat([indices[:, None], boxes], dim=1)
 
 
 class ROIPooler(nn.Module):

detectron2/structures/image_list.py

@@ -5,18 +5,7 @@
 from torch import device
 from torch.nn import functional as F
 
-
-def _as_tensor(x: Tuple[int, int]) -> torch.Tensor:
-    """
-    An equivalent of `torch.as_tensor`, but works under tracing if input
-    is a list of tensor. `torch.as_tensor` will record a constant in tracing,
-    but this function will use `torch.stack` instead.
-    """
-    if torch.jit.is_scripting():
-        return torch.as_tensor(x)
-    if isinstance(x, (list, tuple)) and all([isinstance(t, torch.Tensor) for t in x]):
-        return torch.stack(x)
-    return torch.as_tensor(x)
+from detectron2.layers.wrappers import shapes_to_tensor
 
 
 class ImageList(object):
@@ -90,7 +79,7 @@ def from_tensors(
             assert t.shape[:-2] == tensors[0].shape[:-2], t.shape
 
         image_sizes = [(im.shape[-2], im.shape[-1]) for im in tensors]
-        image_sizes_tensor = [_as_tensor(x) for x in image_sizes]
+        image_sizes_tensor = [shapes_to_tensor(x) for x in image_sizes]
         max_size = torch.stack(image_sizes_tensor).max(0).values
 
         if size_divisibility > 1:

projects/PointRend/point_rend/point_features.py

@@ -2,7 +2,7 @@
 import torch
 from torch.nn import functional as F
 
-from detectron2.layers import cat
+from detectron2.layers import cat, shapes_to_tensor
 from detectron2.structures import BitMasks, Boxes
 
 
@@ -16,15 +16,6 @@
 """
 
 
-def _as_tensor(x):
-    """
-    An equivalent of `torch.as_tensor`, but works under tracing.
-    """
-    if isinstance(x, (list, tuple)) and all([isinstance(t, torch.Tensor) for t in x]):
-        return torch.stack(x)
-    return torch.as_tensor(x)
-
-
 def point_sample(input, point_coords, **kwargs):
     """
     A wrapper around :function:`torch.nn.functional.grid_sample` to support 3D point_coords tensors.
@@ -182,7 +173,7 @@ def point_sample_fine_grained_features(features_list, feature_scales, boxes, poi
         point_features_per_image = []
         for idx_feature, feature_map in enumerate(features_list):
             h, w = feature_map.shape[-2:]
-            scale = _as_tensor([w, h]) / feature_scales[idx_feature]
+            scale = shapes_to_tensor([w, h]) / feature_scales[idx_feature]
             point_coords_scaled = point_coords_wrt_image_per_image / scale.to(feature_map.device)
             point_features_per_image.append(
                 point_sample(

tests/modeling/test_roi_pooler.py

@@ -3,7 +3,7 @@
 import unittest
 import torch
 
-from detectron2.modeling.poolers import ROIPooler, _fmt_box_list
+from detectron2.modeling.poolers import ROIPooler
 from detectron2.structures import Boxes, RotatedBoxes
 from detectron2.utils.testing import random_boxes
 
@@ -114,18 +114,6 @@ def test_no_images(self):
         output = pooler.forward(features, [])
         self.assertEqual(output.shape, (0, C, 14, 14))
 
-    def test_fmt_box_list_tracing(self):
-        class Model(torch.nn.Module):
-            def forward(self, box_tensor):
-                return _fmt_box_list(box_tensor, 0)
-
-        with torch.no_grad():
-            func = torch.jit.trace(Model(), torch.ones(10, 4))
-
-            self.assertEqual(func(torch.ones(10, 4)).shape, (10, 5))
-            self.assertEqual(func(torch.ones(5, 4)).shape, (5, 5))
-            self.assertEqual(func(torch.ones(20, 4)).shape, (20, 5))
-
     def test_roi_pooler_tracing(self):
         class Model(torch.nn.Module):
             def __init__(self, roi):