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))
