def test_rpn_scriptability(self):
cfg = get_cfg()
proposal_generator = RPN(cfg, {
"res4": ShapeSpec(channels=1024, stride=16)
}).eval()
num_images = 2
images_tensor = torch.rand(num_images, 30, 40)
image_sizes = [(32, 32), (30, 40)]
images = ImageList(images_tensor, image_sizes)
features = {"res4": torch.rand(num_images, 1024, 1, 2)}
fields = {"proposal_boxes": Boxes, "objectness_logits": torch.Tensor}
proposal_generator_ts = export_torchscript_with_instances(
proposal_generator, fields)
proposals, _ = proposal_generator(images, features)
proposals_ts, _ = proposal_generator_ts(images, features)
for proposal, proposal_ts in zip(proposals, proposals_ts):
self.assertEqual(proposal.image_size, proposal_ts.image_size)
self.assertTrue(
torch.equal(proposal.proposal_boxes.tensor,
proposal_ts.proposal_boxes.tensor))
self.assertTrue(
torch.equal(proposal.objectness_logits,
proposal_ts.objectness_logits))
def _test_retinanet_model(self, config_path):
model = model_zoo.get(config_path, trained=True)
model.eval()
fields = {
"pred_boxes": Boxes,
"scores": Tensor,
"pred_classes": Tensor,
}
script_model = export_torchscript_with_instances(model, fields)
img = get_sample_coco_image()
inputs = [{"image": img}]
with torch.no_grad():
instance = model(inputs)[0]["instances"]
scripted_instance = convert_scripted_instances(script_model(inputs)[0])
scripted_instance = detector_postprocess(scripted_instance, img.shape[1], img.shape[2])
assert_instances_allclose(instance, scripted_instance)
def export_scripting(torch_model):
assert TORCH_VERSION >= (1, 8)
fields = {
"proposal_boxes": Boxes,
"objectness_logits": Tensor,
"pred_boxes": Boxes,
"scores": Tensor,
"pred_classes": Tensor,
"pred_masks": Tensor,
"pred_keypoints": torch.Tensor,
"pred_keypoint_heatmaps": torch.Tensor,
}
# maybe can export to onnx format?
assert args.format == "torchscript", "Scripting only supports torchscript format."
ts_model = export_torchscript_with_instances(torch_model, fields)
ts_model.save(os.path.join(args.output, "model.ts"))
dump_torchscript_IR(ts_model, args.output)
# TODO inference in Python now missing postprocessing glue code
return None
def _test_rcnn_model(self, config_path):
model = model_zoo.get(config_path, trained=True)
model.eval()
fields = {
"proposal_boxes": Boxes,
"objectness_logits": Tensor,
"pred_boxes": Boxes,
"scores": Tensor,
"pred_classes": Tensor,
"pred_masks": Tensor,
}
script_model = export_torchscript_with_instances(model, fields)
inputs = [{"image": get_sample_coco_image()}]
with torch.no_grad():
instance = model.inference(inputs, do_postprocess=False)[0]
scripted_instance = script_model.inference(inputs, do_postprocess=False)[0]
assert_instances_allclose(instance, scripted_instance)
