def test_script_new_fields(self): class f(torch.nn.Module): def forward(self, x: Instances): proposal_boxes = x.proposal_boxes # noqa F841 objectness_logits = x.objectness_logits # noqa F841 return x class g(torch.nn.Module): def forward(self, x: Instances): mask = x.mask # noqa F841 return x class g2(torch.nn.Module): def forward(self, x: Instances): proposal_boxes = x.proposal_boxes # noqa F841 return x fields = {"proposal_boxes": "Boxes", "objectness_logits": "Tensor"} with patch_instances(fields): torch.jit.script(f()) # can't script anymore after exiting the context with self.assertRaises(Exception): torch.jit.script(g2()) new_fields = {"mask": "Tensor"} with patch_instances(new_fields): torch.jit.script(g()) with self.assertRaises(Exception): torch.jit.script(g2())
def test_script_len(self): class f(torch.nn.Module): def forward(self, x: Instances): return len(x) class g(torch.nn.Module): def forward(self, x: Instances): return len(x) image_shape = (15, 15) fields = {"proposal_boxes": Boxes} with patch_instances(fields) as new_instance: script_module = torch.jit.script(f()) x = new_instance(image_shape) with self.assertRaises(Exception): script_module(x) box_tensors = torch.tensor([[5, 5, 10, 10], [1, 1, 2, 3]]) x.proposal_boxes = Boxes(box_tensors) length = script_module(x) self.assertEqual(length, 2) fields = {"objectness_logits": Tensor} with patch_instances(fields) as new_instance: script_module = torch.jit.script(g()) x = new_instance(image_shape) objectness_logits = torch.tensor([1.0]).reshape(1, 1) x.objectness_logits = objectness_logits length = script_module(x) self.assertEqual(length, 1)
def test_script_new_fields(self): def get_mask(x: Instances) -> torch.Tensor: return x.mask class f(torch.nn.Module): def forward(self, x: Instances): proposal_boxes = x.proposal_boxes # noqa F841 objectness_logits = x.objectness_logits # noqa F841 return x class g(torch.nn.Module): def forward(self, x: Instances): return get_mask(x) class g2(torch.nn.Module): def __init__(self): super().__init__() self.g = g() def forward(self, x: Instances): proposal_boxes = x.proposal_boxes # noqa F841 return x, self.g(x) fields = {"proposal_boxes": Boxes, "objectness_logits": Tensor} with patch_instances(fields): torch.jit.script(f()) # can't script anymore after exiting the context with self.assertRaises(Exception): # will create a ConcreteType for g torch.jit.script(g2()) new_fields = {"mask": Tensor} with patch_instances(new_fields): # will compile g with a different Instances; this should pass torch.jit.script(g()) with self.assertRaises(Exception): torch.jit.script(g2()) new_fields = {"mask": Tensor, "proposal_boxes": Boxes} with patch_instances(new_fields) as NewInstances: # get_mask will be compiled with a different Instances; this should pass scripted_g2 = torch.jit.script(g2()) x = NewInstances((3, 4)) x.mask = torch.rand(3) x.proposal_boxes = Boxes(torch.rand(3, 4)) scripted_g2( x ) # it should accept the new Instances object and run successfully
def test_mask_head_scriptability(self): input_shape = ShapeSpec(channels=1024) mask_features = torch.randn(4, 1024, 14, 14) image_shapes = [(10, 10), (15, 15)] pred_instance0 = Instances(image_shapes[0]) pred_classes0 = torch.tensor([1, 2, 3], dtype=torch.int64) pred_instance0.pred_classes = pred_classes0 pred_instance1 = Instances(image_shapes[1]) pred_classes1 = torch.tensor([4], dtype=torch.int64) pred_instance1.pred_classes = pred_classes1 mask_head = MaskRCNNConvUpsampleHead(input_shape, num_classes=80, conv_dims=[256, 256]).eval() # pred_instance will be in-place changed during the inference # process of `MaskRCNNConvUpsampleHead` origin_outputs = mask_head(mask_features, deepcopy([pred_instance0, pred_instance1])) fields = {"pred_masks": torch.Tensor, "pred_classes": torch.Tensor} with freeze_training_mode(mask_head), patch_instances( fields) as NewInstances: sciript_mask_head = torch.jit.script(mask_head) pred_instance0 = NewInstances.from_instances(pred_instance0) pred_instance1 = NewInstances.from_instances(pred_instance1) script_outputs = sciript_mask_head( mask_features, [pred_instance0, pred_instance1]) for origin_ins, script_ins in zip(origin_outputs, script_outputs): assert_instances_allclose(origin_ins, script_ins.to_instances(), rtol=0)
def test_StandardROIHeads_scriptability(self): cfg = get_cfg() cfg.MODEL.ROI_BOX_HEAD.NAME = "FastRCNNConvFCHead" cfg.MODEL.ROI_BOX_HEAD.NUM_FC = 2 cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE = "ROIAlignV2" cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_WEIGHTS = (10, 10, 5, 5) cfg.MODEL.MASK_ON = True cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = 0.01 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.01 num_images = 2 images_tensor = torch.rand(num_images, 20, 30) image_sizes = [(10, 10), (20, 30)] images = ImageList(images_tensor, image_sizes) num_channels = 1024 features = {"res4": torch.rand(num_images, num_channels, 1, 2)} feature_shape = {"res4": ShapeSpec(channels=num_channels, stride=16)} roi_heads = StandardROIHeads(cfg, feature_shape).eval() proposal0 = Instances(image_sizes[0]) proposal_boxes0 = torch.tensor([[1, 1, 3, 3], [2, 2, 6, 6]], dtype=torch.float32) proposal0.proposal_boxes = Boxes(proposal_boxes0) proposal0.objectness_logits = torch.tensor([0.5, 0.7], dtype=torch.float32) proposal1 = Instances(image_sizes[1]) proposal_boxes1 = torch.tensor([[1, 5, 2, 8], [7, 3, 10, 5]], dtype=torch.float32) proposal1.proposal_boxes = Boxes(proposal_boxes1) proposal1.objectness_logits = torch.tensor([0.1, 0.9], dtype=torch.float32) proposals = [proposal0, proposal1] pred_instances, _ = roi_heads(images, features, proposals) fields = { "objectness_logits": torch.Tensor, "proposal_boxes": Boxes, "pred_classes": torch.Tensor, "scores": torch.Tensor, "pred_masks": torch.Tensor, "pred_boxes": Boxes, "pred_keypoints": torch.Tensor, "pred_keypoint_heatmaps": torch.Tensor, } with freeze_training_mode(roi_heads), patch_instances( fields) as new_instances: proposal0 = new_instances.from_instances(proposal0) proposal1 = new_instances.from_instances(proposal1) proposals = [proposal0, proposal1] scripted_rot_heads = torch.jit.script(roi_heads) scripted_pred_instances, _ = scripted_rot_heads( images, features, proposals) for instance, scripted_instance in zip(pred_instances, scripted_pred_instances): assert_instances_allclose(instance, scripted_instance.to_instances(), rtol=0)
def test_script_access_fields(self): class f(torch.nn.Module): def forward(self, x: Instances): proposal_boxes = x.proposal_boxes objectness_logits = x.objectness_logits return proposal_boxes.tensor + objectness_logits fields = {"proposal_boxes": "Boxes", "objectness_logits": "Tensor"} with patch_instances(fields): torch.jit.script(f())
def test_from_to_instances(self): orig = Instances((30, 30)) orig.proposal_boxes = Boxes(torch.rand(3, 4)) fields = {"proposal_boxes": Boxes, "a": Tensor} with patch_instances(fields) as NewInstances: # convert to NewInstances and back new1 = NewInstances.from_instances(orig) new2 = convert_scripted_instances(new1) 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_has(self): class f(torch.nn.Module): def forward(self, x: Instances): return x.has("proposal_boxes") image_shape = (15, 15) fields = {"proposal_boxes": Boxes} with patch_instances(fields) as new_instance: script_module = torch.jit.script(f()) x = new_instance(image_shape) self.assertFalse(script_module(x)) box_tensors = torch.tensor([[5, 5, 10, 10], [1, 1, 2, 3]]) x.proposal_boxes = Boxes(box_tensors) self.assertTrue(script_module(x))
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_to(self): class f(torch.nn.Module): def forward(self, x: Instances): return x.to(torch.device("cpu")) image_shape = (15, 15) fields = {"proposal_boxes": Boxes, "a": Tensor} with patch_instances(fields) as new_instance: script_module = torch.jit.script(f()) x = new_instance(image_shape) script_module(x) box_tensors = torch.tensor([[5, 5, 10, 10], [1, 1, 2, 3]]) x.proposal_boxes = Boxes(box_tensors) x.a = box_tensors script_module(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])))
def test_script_getitem(self): class f(torch.nn.Module): def forward(self, x: Instances, idx): return x[idx] image_shape = (15, 15) fields = {"proposal_boxes": Boxes, "a": Tensor} inst = Instances(image_shape) inst.proposal_boxes = Boxes(torch.rand(4, 4)) inst.a = torch.rand(4, 10) idx = torch.tensor([True, False, True, False]) with patch_instances(fields) as new_instance: script_module = torch.jit.script(f()) out = f()(inst, idx) out_scripted = script_module(new_instance.from_instances(inst), idx) self.assertTrue( torch.equal(out.proposal_boxes.tensor, out_scripted.proposal_boxes.tensor) ) self.assertTrue(torch.equal(out.a, out_scripted.a))
def test_keypoint_head_scriptability(self): input_shape = ShapeSpec(channels=1024, height=14, width=14) keypoint_features = torch.randn(4, 1024, 14, 14) image_shapes = [(10, 10), (15, 15)] pred_boxes0 = torch.tensor([[1, 1, 3, 3], [2, 2, 6, 6], [1, 5, 2, 8]], dtype=torch.float32) pred_instance0 = Instances(image_shapes[0]) pred_instance0.pred_boxes = Boxes(pred_boxes0) pred_boxes1 = torch.tensor([[7, 3, 10, 5]], dtype=torch.float32) pred_instance1 = Instances(image_shapes[1]) pred_instance1.pred_boxes = Boxes(pred_boxes1) keypoint_head = KRCNNConvDeconvUpsampleHead(input_shape, num_keypoints=17, conv_dims=[512, 512]).eval() origin_outputs = keypoint_head( keypoint_features, deepcopy([pred_instance0, pred_instance1])) fields = { "pred_boxes": "Boxes", "pred_keypoints": "Tensor", "pred_keypoint_heatmaps": "Tensor", } with patch_instances(fields) as NewInstances: sciript_keypoint_head = torch.jit.script(keypoint_head) pred_instance0 = NewInstances.from_instances(pred_instance0) pred_instance1 = NewInstances.from_instances(pred_instance1) script_outputs = sciript_keypoint_head( keypoint_features, [pred_instance0, pred_instance1]) for origin_ins, script_ins in zip(origin_outputs, script_outputs): self.assertEqual(origin_ins.image_size, script_ins.image_size) self.assertTrue( torch.equal(origin_ins.pred_keypoints, script_ins.pred_keypoints)) self.assertTrue( torch.equal(origin_ins.pred_keypoint_heatmaps, script_ins.pred_keypoint_heatmaps))