scripting support for Instances.cat and Instances.__init__(with arguments)

Summary: needed for single stage detector refactoring

Reviewed By: newstzpz

Differential Revision: D30977541

fbshipit-source-id: b596a9996d98caf538a22c8e2e191d4039d013b6

Author: ppwwyyxx

detectron2/export/torchscript_patch.py

@@ -113,18 +113,19 @@ def indent(level, s):
     cls_name = "ScriptedInstances{}".format(_counter)
 
     field_names = tuple(x.name for x in fields)
+    extra_args = ", ".join([f"{f.name}: Optional[{f.annotation}] = None" for f in fields])
     lines.append(
         f"""
 class {cls_name}:
-    def __init__(self, image_size: Tuple[int, int]):
+    def __init__(self, image_size: Tuple[int, int], {extra_args}):
         self.image_size = image_size
         self._field_names = {field_names}
 """
     )
 
     for f in fields:
         lines.append(
-            indent(2, f"self._{f.name} = torch.jit.annotate(Optional[{f.annotation}], None)")
+            indent(2, f"self._{f.name} = torch.jit.annotate(Optional[{f.annotation}], {f.name})")
         )
 
     for f in fields:
@@ -135,7 +136,7 @@ def {f.name}(self) -> {f.annotation}:
         # has to use a local for type refinement
         # https://pytorch.org/docs/stable/jit_language_reference.html#optional-type-refinement
         t = self._{f.name}
-        assert t is not None
+        assert t is not None, "{f.name} is None and cannot be accessed!"
         return t
 
     @{f.name}.setter
@@ -184,10 +185,11 @@ def has(self, name: str) -> bool:
     )
 
     # support method `to`
+    none_args = ", None" * len(fields)
     lines.append(
         f"""
     def to(self, device: torch.device) -> "{cls_name}":
-        ret = {cls_name}(self.image_size)
+        ret = {cls_name}(self.image_size{none_args})
 """
     )
     for f in fields:
@@ -210,10 +212,11 @@ def to(self, device: torch.device) -> "{cls_name}":
     )
 
     # support method `getitem`
+    none_args = ", None" * len(fields)
     lines.append(
         f"""
     def __getitem__(self, item) -> "{cls_name}":
-        ret = {cls_name}(self.image_size)
+        ret = {cls_name}(self.image_size{none_args})
 """
     )
     for f in fields:
@@ -230,6 +233,32 @@ def __getitem__(self, item) -> "{cls_name}":
 """
     )
 
+    # support method `cat`
+    # this version does not contain checks that all instances have same size and fields
+    none_args = ", None" * len(fields)
+    lines.append(
+        f"""
+    def cat(self, instances: List["{cls_name}"]) -> "{cls_name}":
+        ret = {cls_name}(self.image_size{none_args})
+"""
+    )
+    for f in fields:
+        lines.append(
+            f"""
+        t = self._{f.name}
+        if t is not None:
+            values: List[{f.annotation}] = [x.{f.name} for x in instances]
+            if torch.jit.isinstance(t, torch.Tensor):
+                ret._{f.name} = torch.cat(values, dim=0)
+            else:
+                ret._{f.name} = t.cat(values)
+"""
+        )
+    lines.append(
+        """
+        return ret"""
+    )
+
     # support method `get_fields()`
     lines.append(
         """

tests/structures/test_instances.py

@@ -181,6 +181,41 @@ def test_from_to_instances(self):
         self.assertTrue(torch.equal(orig.proposal_boxes.tensor, new1.proposal_boxes.tensor))
         self.assertTrue(torch.equal(orig.proposal_boxes.tensor, new2.proposal_boxes.tensor))
 
+    def test_script_init_args(self):
+        def f(x: Tensor):
+            image_shape = (15, 15)
+            # __init__ can take arguments
+            inst = Instances(image_shape, a=x, proposal_boxes=Boxes(x))
+            inst2 = Instances(image_shape, a=x)
+            return inst.a, inst2.a
+
+        fields = {"proposal_boxes": Boxes, "a": Tensor}
+        with patch_instances(fields):
+            script_f = torch.jit.script(f)
+            x = torch.randn(3, 4)
+            outputs = script_f(x)
+            self.assertTrue(torch.equal(outputs[0], x))
+            self.assertTrue(torch.equal(outputs[1], x))
+
+    def test_script_cat(self):
+        def f(x: Tensor):
+            image_shape = (15, 15)
+            # __init__ can take arguments
+            inst = Instances(image_shape, a=x)
+            inst2 = Instances(image_shape, a=x)
+
+            inst3 = Instances(image_shape, proposal_boxes=Boxes(x))
+            return inst.cat([inst, inst2]), inst3.cat([inst3, inst3])
+
+        fields = {"proposal_boxes": Boxes, "a": Tensor}
+        with patch_instances(fields):
+            script_f = torch.jit.script(f)
+            x = torch.randn(3, 4)
+            output, output2 = script_f(x)
+            self.assertTrue(torch.equal(output.a, torch.cat([x, x])))
+            self.assertFalse(output.has("proposal_boxes"))
+            self.assertTrue(torch.equal(output2.proposal_boxes.tensor, torch.cat([x, x])))
+
 
 if __name__ == "__main__":
     unittest.main()