Add YOLO object detection model

CHANGELOG.md

@@ -88,6 +88,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Added Pix2Pix model ([#533](https://github.com/PyTorchLightning/lightning-bolts/pull/533))
+- Added YOLO model ([#552](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/552))
 
 ### Changed
 

docs/source/index.rst

@@ -71,6 +71,7 @@ Lightning-Bolts documentation
 
    autoencoders
    convolutional
+   object_detection
    gans
    reinforce_learn
    self_supervised_models

docs/source/object_detection.rst

@@ -0,0 +1,20 @@
+Object Detection
+================
+This package lists contributed object detection models.
+
+--------------
+
+
+Faster R-CNN
+------------
+
+.. autoclass:: pl_bolts.models.detection.faster_rcnn.faster_rcnn_module.FasterRCNN
+    :noindex:
+
+-------------
+
+YOLO
+----
+
+.. autoclass:: pl_bolts.models.detection.yolo.yolo_module.YOLO
+    :noindex:

pl_bolts/datamodules/vocdetection_datamodule.py

@@ -1,9 +1,10 @@
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import os
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
 from pytorch_lightning import LightningDataModule
 from torch import Tensor
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 
 from pl_bolts.utils import _TORCHVISION_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
@@ -107,10 +108,11 @@ class VOCDetectionDataModule(LightningDataModule):
 
     def __init__(
         self,
-        data_dir: str,
+        data_dir: Optional[str] = None,
         year: str = "2012",
         num_workers: int = 0,
         normalize: bool = False,
+        batch_size: int = 16,
         shuffle: bool = True,
         pin_memory: bool = True,
         drop_last: bool = False,
@@ -125,9 +127,10 @@ def __init__(
         super().__init__(*args, **kwargs)
 
         self.year = year
-        self.data_dir = data_dir
+        self.data_dir = data_dir if data_dir is not None else os.getcwd()
         self.num_workers = num_workers
         self.normalize = normalize
+        self.batch_size = batch_size
         self.shuffle = shuffle
         self.pin_memory = pin_memory
         self.drop_last = drop_last
@@ -145,60 +148,50 @@ def prepare_data(self) -> None:
         VOCDetection(self.data_dir, year=self.year, image_set="train", download=True)
         VOCDetection(self.data_dir, year=self.year, image_set="val", download=True)
 
-    def train_dataloader(
-        self, batch_size: int = 1, image_transforms: Union[List[Callable], Callable] = None
-    ) -> DataLoader:
+    def train_dataloader(self, image_transforms: Optional[Callable] = None) -> DataLoader:
         """VOCDetection train set uses the `train` subset.
 
         Args:
-            batch_size: size of batch
-            transforms: custom transforms
+            image_transforms: custom image-only transforms
         """
-        transforms = [_prepare_voc_instance]
-        image_transforms = image_transforms or self.train_transforms or self._default_transforms()
+        transforms = [
+            _prepare_voc_instance,
+            self.default_transforms() if self.train_transforms is None else self.train_transforms,
+        ]
         transforms = Compose(transforms, image_transforms)
+
         dataset = VOCDetection(self.data_dir, year=self.year, image_set="train", transforms=transforms)
-        loader = DataLoader(
-            dataset,
-            batch_size=batch_size,
-            shuffle=self.shuffle,
-            num_workers=self.num_workers,
-            drop_last=self.drop_last,
-            pin_memory=self.pin_memory,
-            collate_fn=_collate_fn,
-        )
-        return loader
+        return self._data_loader(dataset, shuffle=self.shuffle)
 
-    def val_dataloader(self, batch_size: int = 1, image_transforms: Optional[List[Callable]] = None) -> DataLoader:
+    def val_dataloader(self, image_transforms: Optional[Callable] = None) -> DataLoader:
         """VOCDetection val set uses the `val` subset.
 
         Args:
-            batch_size: size of batch
-            transforms: custom transforms
+            image_transforms: custom image-only transforms
         """
-        transforms = [_prepare_voc_instance]
-        image_transforms = image_transforms or self.train_transforms or self._default_transforms()
+        transforms = [
+            _prepare_voc_instance,
+            self.default_transforms() if self.val_transforms is None else self.val_transforms,
+        ]
         transforms = Compose(transforms, image_transforms)
+
         dataset = VOCDetection(self.data_dir, year=self.year, image_set="val", transforms=transforms)
-        loader = DataLoader(
+        return self._data_loader(dataset, shuffle=False)
+
+    def default_transforms(self) -> Callable:
+        voc_transforms = [transform_lib.ToTensor()]
+        if self.normalize:
+            voc_transforms += [transform_lib.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
+        voc_transforms = transform_lib.Compose(voc_transforms)
+        return lambda image, target: (voc_transforms(image), target)
+
+    def _data_loader(self, dataset: Dataset, shuffle: bool = False) -> DataLoader:
+        return DataLoader(
             dataset,
-            batch_size=batch_size,
-            shuffle=False,
+            batch_size=self.batch_size,
+            shuffle=shuffle,
             num_workers=self.num_workers,
             drop_last=self.drop_last,
             pin_memory=self.pin_memory,
             collate_fn=_collate_fn,
         )
-        return loader
-
-    def _default_transforms(self) -> Callable:
-        if self.normalize:
-            voc_transforms = transform_lib.Compose(
-                [
-                    transform_lib.ToTensor(),
-                    transform_lib.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-                ]
-            )
-        else:
-            voc_transforms = transform_lib.Compose([transform_lib.ToTensor()])
-        return voc_transforms

pl_bolts/models/detection/__init__.py

@@ -1,7 +1,6 @@
 from pl_bolts.models.detection import components
 from pl_bolts.models.detection.faster_rcnn import FasterRCNN
+from pl_bolts.models.detection.yolo.yolo_config import YOLOConfiguration
+from pl_bolts.models.detection.yolo.yolo_module import YOLO
 
-__all__ = [
-    "components",
-    "FasterRCNN",
-]
+__all__ = ["components", "FasterRCNN", "YOLOConfiguration", "YOLO"]

pl_bolts/models/detection/faster_rcnn/faster_rcnn_module.py

@@ -147,9 +147,8 @@ def run_cli():
 
     seed_everything(42)
     parser = ArgumentParser()
+    parser = VOCDetectionDataModule.add_argparse_args(parser)
     parser = Trainer.add_argparse_args(parser)
-    parser.add_argument("--data_dir", type=str, default=".")
-    parser.add_argument("--batch_size", type=int, default=1)
     parser = FasterRCNN.add_model_specific_args(parser)
 
     args = parser.parse_args()