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 = scripting_with_instances(model, fields) inputs = [{"image": get_sample_coco_image()}] * 2 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)
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 = scripting_with_instances(model, fields) img = get_sample_coco_image() inputs = [{"image": img}] * 2 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 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 = scripting_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_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 = scripting_with_instances(model, fields) # Test that batch inference with different shapes are supported image = get_sample_coco_image() small_image = nn.functional.interpolate(image, scale_factor=0.5) inputs = [{"image": image}, {"image": small_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)
import torch from detectron2.config import get_cfg from detectron2.export import scripting_with_instances from detectron2.modeling import build_model from detectron2.structures import Boxes cfg = get_cfg() cfg.merge_from_file("configs/faster_rcnn_R_50_FPN_1x.yaml") cfg.MODEL.WEIGHTS = "model/model_a.pth" cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = ( 1 ) cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 5 cfg.MODEL.DEVICE = "cpu" model = build_model(cfg) model.eval() fields = {"pred_boxes": Boxes, "scores": torch.Tensor, "pred_classes": torch.Tensor, "proposal_boxes": Boxes, "objectness_logits": torch.Tensor} torchscipt_model = scripting_with_instances(model, fields) torch.jit.save(torchscipt_model, "model/jit_a.pth") # 提取缺陷