Add RetinaNet Object detection with Backbones

pl_bolts/metrics/object_detection.py

@@ -1,6 +1,18 @@
 import torch
 
 
+def _evaluate_iou(preds: torch.Tensor, target: torch.Tensor):
+    """
+    Evaluate intersection over union (IOU) for target from dataset and output prediction
+    from model
+    """
+
+    if preds["boxes"].shape[0] == 0:
+        # no box detected, 0 IOU
+        return torch.tensor(0.0, device=preds["boxes"].device)
+    return iou(target["boxes"], preds["boxes"]).diag().mean()
+
+
 def iou(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
     """
     Calculates the intersection over union.
@@ -58,6 +70,12 @@ def giou(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         GIoU in an NxM tensor containing the pairwise GIoU values for every element in preds and target,
         where N is the number of prediction bounding boxes and M is the number of target bounding boxes
     """
+
+    # degenerate boxes gives inf / nan results
+    # so do an early check
+    assert (preds[:, 2:] >= preds[:, :2]).all()
+    assert (target[:, 2:] >= target[:, :2]).all()
+
     x_min = torch.max(preds[:, None, 0], target[:, 0])
     y_min = torch.max(preds[:, None, 1], target[:, 1])
     x_max = torch.min(preds[:, None, 2], target[:, 2])

pl_bolts/models/detection/__init__.py

@@ -1,7 +1,5 @@
 from pl_bolts.models.detection import components  # noqa: F401
 from pl_bolts.models.detection.faster_rcnn import FasterRCNN  # noqa: F401
+from pl_bolts.models.detection.retinanet import RetinaNet  # noqa: F401
 
-__all__ = [
-    "components",
-    "FasterRCNN",
-]
+__all__ = ["components", "FasterRCNN", "RetinaNet"]

pl_bolts/models/detection/faster_rcnn/faster_rcnn_module.py

@@ -4,32 +4,18 @@
 import pytorch_lightning as pl
 import torch
 
+from pl_bolts.metrics.object_detection import _evaluate_iou
 from pl_bolts.models.detection.faster_rcnn import create_fasterrcnn_backbone
 from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
 if _TORCHVISION_AVAILABLE:
     from torchvision.models.detection.faster_rcnn import FasterRCNN as torchvision_FasterRCNN
     from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn, FastRCNNPredictor
-    from torchvision.ops import box_iou
 else:  # pragma: no cover
     warn_missing_pkg("torchvision")
 
 
-def _evaluate_iou(target, pred):
-    """
-    Evaluate intersection over union (IOU) for target from dataset and output prediction
-    from model
-    """
-    if not _TORCHVISION_AVAILABLE:  # pragma: no cover
-        raise ModuleNotFoundError('You want to use `torchvision` which is not installed yet.')
-
-    if pred["boxes"].shape[0] == 0:
-        # no box detected, 0 IOU
-        return torch.tensor(0.0, device=pred["boxes"].device)
-    return box_iou(target["boxes"], pred["boxes"]).diag().mean()
-
-
 class FasterRCNN(pl.LightningModule):
     """
     PyTorch Lightning implementation of `Faster R-CNN: Towards Real-Time Object Detection with
@@ -47,7 +33,7 @@ class FasterRCNN(pl.LightningModule):
     CLI command::
 
         # PascalVOC
-        python faster_rcnn.py --gpus 1 --pretrained True
+        python faster_rcnn_module.py --gpus 1 --pretrained True
     """
 
     def __init__(
@@ -117,7 +103,7 @@ def validation_step(self, batch, batch_idx):
         images, targets = batch
         # fasterrcnn takes only images for eval() mode
         outs = self.model(images)
-        iou = torch.stack([_evaluate_iou(t, o) for t, o in zip(targets, outs)]).mean()
+        iou = torch.stack([_evaluate_iou(o, t) for t, o in zip(targets, outs)]).mean()
         return {"val_iou": iou}
 
     def validation_epoch_end(self, outs):

pl_bolts/models/detection/retinanet/__init__.py

@@ -0,0 +1,4 @@
+from pl_bolts.models.detection.retinanet.backbones import create_retinanet_backbone
+from pl_bolts.models.detection.retinanet.retainanet_module import RetinaNet
+
+__all__ = ["create_retinanet_backbone", "RetinaNet"]

pl_bolts/models/detection/retinanet/backbones.py

@@ -0,0 +1,41 @@
+from typing import Any, Optional
+
+import torch.nn as nn
+
+from pl_bolts.models.detection.components import create_torchvision_backbone
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
+from pl_bolts.utils.warnings import warn_missing_pkg
+
+if _TORCHVISION_AVAILABLE:
+    from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
+else:  # pragma: no cover
+    warn_missing_pkg("torchvision")
+
+
+def create_retinanet_backbone(
+    backbone: str,
+    fpn: bool = True,
+    pretrained: Optional[str] = None,
+    trainable_backbone_layers: int = 3,
+    **kwargs: Any
+) -> nn.Module:
+    """
+    Args:
+        backbone:
+            Supported backones are: "resnet18", "resnet34","resnet50", "resnet101", "resnet152",
+            "resnext50_32x4d", "resnext101_32x8d", "wide_resnet50_2", "wide_resnet101_2",
+            as resnets with fpn backbones.
+            Without fpn backbones supported are: "resnet18", "resnet34", "resnet50","resnet101",
+            "resnet152", "resnext101_32x8d", "mobilenet_v2", "vgg11", "vgg13", "vgg16", "vgg19",
+        fpn: If True then constructs fpn as well.
+        pretrained: If None creates imagenet weights backbone.
+        trainable_backbone_layers: number of trainable resnet layers starting from final block.
+    """
+
+    if fpn:
+        # Creates a torchvision resnet model with fpn added.
+        backbone = resnet_fpn_backbone(backbone, pretrained=True, trainable_layers=trainable_backbone_layers, **kwargs)
+    else:
+        # This does not create fpn backbone, it is supported for all models
+        backbone, _ = create_torchvision_backbone(backbone, pretrained)
+    return backbone

pl_bolts/models/detection/retinanet/retainanet_module.py

@@ -0,0 +1,148 @@
+from argparse import ArgumentParser
+from typing import Any, Optional
+
+import pytorch_lightning as pl
+import torch
+
+from pl_bolts.metrics.object_detection import _evaluate_iou
+from pl_bolts.models.detection.retinanet import create_retinanet_backbone
+from pl_bolts.utils import _TORCHVISION_AVAILABLE
+from pl_bolts.utils.warnings import warn_missing_pkg
+
+if _TORCHVISION_AVAILABLE:
+    from torchvision.models.detection.retinanet import RetinaNet as torchvision_RetinaNet
+    from torchvision.models.detection.retinanet import retinanet_resnet50_fpn, RetinaNetHead
+else:  # pragma: no cover
+    warn_missing_pkg("torchvision")
+
+
+class RetinaNet(pl.LightningModule):
+    """
+    PyTorch Lightning implementation of Retina Net `Focal Loss for
+    Dense Object Detection <https://arxiv.org/abs/1708.02002>`_.
+
+    Paper authors: Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, Piotr Dollár
+
+    Model implemented by:
+        - `Aditya Oke <https://github.com/oke-aditya>`
+
+    During training, the model expects both the input tensors, as well as targets (list of dictionary), containing:
+        - boxes (`FloatTensor[N, 4]`): the ground truth boxes in `[x1, y1, x2, y2]` format.
+        - labels (`Int64Tensor[N]`): the class label for each ground truh box
+
+    CLI command::
+
+        # PascalVOC
+        python retinanet_module.py --gpus 1 --pretrained True
+    """
+
+    def __init__(
+        self,
+        learning_rate: float = 0.0001,
+        num_classes: int = 91,
+        backbone: Optional[str] = None,
+        fpn: bool = True,
+        pretrained: bool = False,
+        pretrained_backbone: bool = True,
+        trainable_backbone_layers: int = 3,
+        **kwargs: Any,
+    ):
+        """
+        Args:
+            learning_rate: the learning rate
+            num_classes: number of detection classes (including background)
+            backbone: Pretained backbone CNN architecture.
+            fpn: If True, creates a Feature Pyramind Network on top of Resnet based CNNs.
+            pretrained: if true, returns a model pre-trained on COCO train2017
+            pretrained_backbone: if true, returns a model with backbone pre-trained on Imagenet
+            trainable_backbone_layers: number of trainable resnet layers starting from final block
+        """
+        super().__init__()
+        self.learning_rate = learning_rate
+        self.num_classes = num_classes
+        self.backbone = backbone
+        if backbone is None:
+            self.model = retinanet_resnet50_fpn(pretrained=pretrained, **kwargs)
+
+            self.model.head = RetinaNetHead(
+                in_channels=self.model.backbone.out_channels,
+                num_anchors=self.model.head.classification_head.num_anchors,
+                num_classes=num_classes,
+                **kwargs
+            )
+
+        else:
+            backbone_model = create_retinanet_backbone(
+                self.backbone, fpn, pretrained_backbone, trainable_backbone_layers, **kwargs
+            )
+            self.model = torchvision_RetinaNet(backbone_model, num_classes=num_classes, **kwargs)
+
+    def forward(self, x):
+        self.model.eval()
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+
+        images, targets = batch
+        targets = [{k: v for k, v in t.items()} for t in targets]
+
+        # fasterrcnn takes both images and targets for training, returns
+        loss_dict = self.model(images, targets)
+        loss = sum(loss for loss in loss_dict.values())
+        return {"loss": loss, "log": loss_dict}
+
+    def validation_step(self, batch, batch_idx):
+        images, targets = batch
+        # fasterrcnn takes only images for eval() mode
+        outs = self.model(images)
+        iou = torch.stack([_evaluate_iou(o, t) for t, o in zip(targets, outs)]).mean()
+        return {"val_iou": iou}
+
+    def validation_epoch_end(self, outs):
+        avg_iou = torch.stack([o["val_iou"] for o in outs]).mean()
+        logs = {"val_iou": avg_iou}
+        return {"avg_val_iou": avg_iou, "log": logs}
+
+    def configure_optimizers(self):
+        return torch.optim.SGD(
+            self.model.parameters(),
+            lr=self.learning_rate,
+            momentum=0.9,
+            weight_decay=0.005,
+        )
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+        parser.add_argument("--learning_rate", type=float, default=0.0001)
+        parser.add_argument("--num_classes", type=int, default=91)
+        parser.add_argument("--backbone", type=str, default=None)
+        parser.add_argument("--fpn", type=bool, default=True)
+        parser.add_argument("--pretrained", type=bool, default=False)
+        parser.add_argument("--pretrained_backbone", type=bool, default=True)
+        parser.add_argument("--trainable_backbone_layers", type=int, default=3)
+        return parser
+
+
+def run_cli():
+    from pl_bolts.datamodules import VOCDetectionDataModule
+
+    pl.seed_everything(42)
+    parser = ArgumentParser()
+    parser = pl.Trainer.add_argparse_args(parser)
+    parser.add_argument("--data_dir", type=str, default=".")
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser = RetinaNet.add_model_specific_args(parser)
+
+    args = parser.parse_args()
+
+    datamodule = VOCDetectionDataModule.from_argparse_args(args)
+    args.num_classes = datamodule.num_classes
+
+    model = RetinaNet(**vars(args))
+    trainer = pl.Trainer.from_argparse_args(args)
+    trainer.fit(model, datamodule=datamodule)
+
+
+if __name__ == "__main__":
+    run_cli()

requirements.txt

@@ -1,2 +1,2 @@
-torch>=1.6
+torch>=1.7
 pytorch-lightning>=1.1.1

requirements/models.txt

@@ -1,4 +1,4 @@
-torchvision>=0.7
+torchvision>=0.8.1
 scikit-learn>=0.23
 Pillow
 opencv-python

tests/models/test_detection.py

@@ -3,13 +3,14 @@
 from torch.utils.data import DataLoader
 
 from pl_bolts.datasets import DummyDetectionDataset
-from pl_bolts.models.detection import FasterRCNN
+from pl_bolts.models.detection import FasterRCNN, RetinaNet
 
 
 def _collate_fn(batch):
     return tuple(zip(*batch))
 
 
+@torch.no_grad()
 def test_fasterrcnn():
     model = FasterRCNN()
 
@@ -18,19 +19,54 @@ def test_fasterrcnn():
 
 
 def test_fasterrcnn_train(tmpdir):
-    model = FasterRCNN()
+    model = FasterRCNN(pretrained=False)
 
     train_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
     valid_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
 
-    trainer = pl.Trainer(fast_dev_run=True, default_root_dir=tmpdir)
+    trainer = pl.Trainer(
+        fast_dev_run=True, logger=False, checkpoint_callback=False, max_epochs=1, default_root_dir=tmpdir
+    )
     trainer.fit(model, train_dataloader=train_dl, val_dataloaders=valid_dl)
 
 
 def test_fasterrcnn_bbone_train(tmpdir):
-    model = FasterRCNN(backbone="resnet18", fpn=True, pretrained_backbone=True)
+    model = FasterRCNN(backbone="resnet18", fpn=True, pretrained_backbone=False)
+    train_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
+    valid_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
+
+    trainer = pl.Trainer(
+        fast_dev_run=True, logger=False, checkpoint_callback=False, max_epochs=1, default_root_dir=tmpdir
+    )
+    trainer.fit(model, train_dl, valid_dl)
+
+
+@torch.no_grad()
+def test_retinanet():
+    model = RetinaNet(pretrained=False)
+
+    image = torch.rand(1, 3, 400, 400)
+    model(image)
+
+
+def test_retinanet_train(tmpdir):
+    model = RetinaNet(pretrained=False)
+
+    train_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
+    valid_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
+
+    trainer = pl.Trainer(
+        fast_dev_run=True, logger=False, checkpoint_callback=False, max_epochs=1, default_root_dir=tmpdir
+    )
+    trainer.fit(model, train_dataloader=train_dl, val_dataloaders=valid_dl)
+
+
+def test_retinanet_bbone_train(tmpdir):
+    model = RetinaNet(backbone="resnet18", fpn=True, pretrained_backbone=False)
     train_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
     valid_dl = DataLoader(DummyDetectionDataset(), collate_fn=_collate_fn)
 
-    trainer = pl.Trainer(fast_dev_run=True, default_root_dir=tmpdir)
+    trainer = pl.Trainer(
+        fast_dev_run=True, logger=False, checkpoint_callback=False, max_epochs=1, default_root_dir=tmpdir
+    )
     trainer.fit(model, train_dl, valid_dl)