Spaces:
Running
Running
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import logging | |
| import unittest | |
| import torch | |
| from detectron2.layers import ShapeSpec | |
| from detectron2.modeling.box_regression import Box2BoxTransform, Box2BoxTransformRotated | |
| from detectron2.modeling.roi_heads.fast_rcnn import FastRCNNOutputLayers | |
| from detectron2.modeling.roi_heads.rotated_fast_rcnn import RotatedFastRCNNOutputLayers | |
| from detectron2.structures import Boxes, Instances, RotatedBoxes | |
| from detectron2.utils.events import EventStorage | |
| logger = logging.getLogger(__name__) | |
| class FastRCNNTest(unittest.TestCase): | |
| def test_fast_rcnn(self): | |
| torch.manual_seed(132) | |
| box_head_output_size = 8 | |
| box_predictor = FastRCNNOutputLayers( | |
| ShapeSpec(channels=box_head_output_size), | |
| box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)), | |
| num_classes=5, | |
| ) | |
| feature_pooled = torch.rand(2, box_head_output_size) | |
| predictions = box_predictor(feature_pooled) | |
| proposal_boxes = torch.tensor([[0.8, 1.1, 3.2, 2.8], [2.3, 2.5, 7, 8]], dtype=torch.float32) | |
| gt_boxes = torch.tensor([[1, 1, 3, 3], [2, 2, 6, 6]], dtype=torch.float32) | |
| proposal = Instances((10, 10)) | |
| proposal.proposal_boxes = Boxes(proposal_boxes) | |
| proposal.gt_boxes = Boxes(gt_boxes) | |
| proposal.gt_classes = torch.tensor([1, 2]) | |
| with EventStorage(): # capture events in a new storage to discard them | |
| losses = box_predictor.losses(predictions, [proposal]) | |
| expected_losses = { | |
| "loss_cls": torch.tensor(1.7951188087), | |
| "loss_box_reg": torch.tensor(4.0357131958), | |
| } | |
| for name in expected_losses.keys(): | |
| assert torch.allclose(losses[name], expected_losses[name]) | |
| def test_fast_rcnn_empty_batch(self, device="cpu"): | |
| box_predictor = FastRCNNOutputLayers( | |
| ShapeSpec(channels=10), | |
| box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)), | |
| num_classes=8, | |
| ).to(device=device) | |
| logits = torch.randn(0, 100, requires_grad=True, device=device) | |
| deltas = torch.randn(0, 4, requires_grad=True, device=device) | |
| losses = box_predictor.losses([logits, deltas], []) | |
| for value in losses.values(): | |
| self.assertTrue(torch.allclose(value, torch.zeros_like(value))) | |
| sum(losses.values()).backward() | |
| self.assertTrue(logits.grad is not None) | |
| self.assertTrue(deltas.grad is not None) | |
| predictions, _ = box_predictor.inference([logits, deltas], []) | |
| self.assertEqual(len(predictions), 0) | |
| def test_fast_rcnn_empty_batch_cuda(self): | |
| self.test_fast_rcnn_empty_batch(device=torch.device("cuda")) | |
| def test_fast_rcnn_rotated(self): | |
| torch.manual_seed(132) | |
| box_head_output_size = 8 | |
| box_predictor = RotatedFastRCNNOutputLayers( | |
| ShapeSpec(channels=box_head_output_size), | |
| box2box_transform=Box2BoxTransformRotated(weights=(10, 10, 5, 5, 1)), | |
| num_classes=5, | |
| ) | |
| feature_pooled = torch.rand(2, box_head_output_size) | |
| predictions = box_predictor(feature_pooled) | |
| proposal_boxes = torch.tensor( | |
| [[2, 1.95, 2.4, 1.7, 0], [4.65, 5.25, 4.7, 5.5, 0]], dtype=torch.float32 | |
| ) | |
| gt_boxes = torch.tensor([[2, 2, 2, 2, 0], [4, 4, 4, 4, 0]], dtype=torch.float32) | |
| proposal = Instances((10, 10)) | |
| proposal.proposal_boxes = RotatedBoxes(proposal_boxes) | |
| proposal.gt_boxes = RotatedBoxes(gt_boxes) | |
| proposal.gt_classes = torch.tensor([1, 2]) | |
| with EventStorage(): # capture events in a new storage to discard them | |
| losses = box_predictor.losses(predictions, [proposal]) | |
| # Note: the expected losses are slightly different even if | |
| # the boxes are essentially the same as in the FastRCNNOutput test, because | |
| # bbox_pred in FastRCNNOutputLayers have different Linear layers/initialization | |
| # between the two cases. | |
| expected_losses = { | |
| "loss_cls": torch.tensor(1.7920907736), | |
| "loss_box_reg": torch.tensor(4.0410838127), | |
| } | |
| for name in expected_losses.keys(): | |
| assert torch.allclose(losses[name], expected_losses[name]) | |
| def test_predict_boxes_tracing(self): | |
| class Model(torch.nn.Module): | |
| def __init__(self, output_layer): | |
| super(Model, self).__init__() | |
| self._output_layer = output_layer | |
| def forward(self, proposal_deltas, proposal_boxes): | |
| instances = Instances((10, 10)) | |
| instances.proposal_boxes = Boxes(proposal_boxes) | |
| return self._output_layer.predict_boxes((None, proposal_deltas), [instances]) | |
| box_head_output_size = 8 | |
| box_predictor = FastRCNNOutputLayers( | |
| ShapeSpec(channels=box_head_output_size), | |
| box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)), | |
| num_classes=5, | |
| ) | |
| model = Model(box_predictor) | |
| from detectron2.export.torchscript_patch import patch_builtin_len | |
| with torch.no_grad(), patch_builtin_len(): | |
| func = torch.jit.trace(model, (torch.randn(10, 20), torch.randn(10, 4))) | |
| o = func(torch.randn(10, 20), torch.randn(10, 4)) | |
| self.assertEqual(o[0].shape, (10, 20)) | |
| o = func(torch.randn(5, 20), torch.randn(5, 4)) | |
| self.assertEqual(o[0].shape, (5, 20)) | |
| o = func(torch.randn(20, 20), torch.randn(20, 4)) | |
| self.assertEqual(o[0].shape, (20, 20)) | |
| def test_predict_probs_tracing(self): | |
| class Model(torch.nn.Module): | |
| def __init__(self, output_layer): | |
| super(Model, self).__init__() | |
| self._output_layer = output_layer | |
| def forward(self, scores, proposal_boxes): | |
| instances = Instances((10, 10)) | |
| instances.proposal_boxes = Boxes(proposal_boxes) | |
| return self._output_layer.predict_probs((scores, None), [instances]) | |
| box_head_output_size = 8 | |
| box_predictor = FastRCNNOutputLayers( | |
| ShapeSpec(channels=box_head_output_size), | |
| box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)), | |
| num_classes=5, | |
| ) | |
| model = Model(box_predictor) | |
| from detectron2.export.torchscript_patch import patch_builtin_len | |
| with torch.no_grad(), patch_builtin_len(): | |
| func = torch.jit.trace(model, (torch.randn(10, 6), torch.rand(10, 4))) | |
| o = func(torch.randn(10, 6), torch.randn(10, 4)) | |
| self.assertEqual(o[0].shape, (10, 6)) | |
| o = func(torch.randn(5, 6), torch.randn(5, 4)) | |
| self.assertEqual(o[0].shape, (5, 6)) | |
| o = func(torch.randn(20, 6), torch.randn(20, 4)) | |
| self.assertEqual(o[0].shape, (20, 6)) | |
| if __name__ == "__main__": | |
| unittest.main() | |