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")
# 提取缺陷
