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)