detectron2/tests/modeling/test_fast_rcnn.py

# 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)

    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
    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()