Add YOLO object detection model

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
     [#323](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/323))
 - Added data monitor callbacks `ModuleDataMonitor` and `TrainingDataMonitor` ([#285](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/285))
 - Added DCGAN module ([#403](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/403)) 
+- Added YOLO module ([#552](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/552))
 - Added `VisionDataModule` as parent class for `BinaryMNISTDataModule`, `CIFAR10DataModule`, `FashionMNISTDataModule`, 
   and `MNISTDataModule` ([#400](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/400))
 - Added GIoU loss ([#347](https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/347))

docs/source/index.rst

@@ -70,6 +70,7 @@ PyTorch-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,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
 from pytorch_lightning import LightningDataModule
@@ -152,16 +152,20 @@ def prepare_data(self) -> None:
         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
+        self,
+        batch_size: int = 1,
+        transforms: List[Callable] = [],
+        image_transforms: Optional[Callable] = None
     ) -> DataLoader:
         """
         VOCDetection train set uses the `train` subset
 
         Args:
             batch_size: size of batch
-            transforms: custom transforms
+            transforms: custom transforms for image and target
+            image_transforms: custom image-only transforms
         """
-        transforms = [_prepare_voc_instance]
+        transforms = [_prepare_voc_instance] + transforms
         image_transforms = image_transforms or self.train_transforms or self._default_transforms()
         transforms = Compose(transforms, image_transforms)
         dataset = VOCDetection(self.data_dir, year=self.year, image_set="train", transforms=transforms)
@@ -176,15 +180,21 @@ def train_dataloader(
         )
         return loader
 
-    def val_dataloader(self, batch_size: int = 1, image_transforms: Optional[List[Callable]] = None) -> DataLoader:
+    def val_dataloader(
+        self,
+        batch_size: int = 1,
+        transforms: List[Callable] = [],
+        image_transforms: Optional[Callable] = None
+    ) -> DataLoader:
         """
         VOCDetection val set uses the `val` subset
 
         Args:
             batch_size: size of batch
-            transforms: custom transforms
+            transforms: custom transforms for image and target
+            image_transforms: custom image-only transforms
         """
-        transforms = [_prepare_voc_instance]
+        transforms = [_prepare_voc_instance] + transforms
         image_transforms = image_transforms or self.train_transforms or self._default_transforms()
         transforms = Compose(transforms, image_transforms)
         dataset = VOCDetection(self.data_dir, year=self.year, image_set="val", transforms=transforms)

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 import components
+from pl_bolts.models.detection.faster_rcnn import FasterRCNN
+from pl_bolts.models.detection.yolo import Yolo, YoloConfiguration
 
-__all__ = [
-    "components",
-    "FasterRCNN",
-]
+__all__ = ["components", "FasterRCNN", "YoloConfiguration", "Yolo"]

pl_bolts/models/detection/yolo/__init__.py

@@ -0,0 +1,4 @@
+from pl_bolts.models.detection.yolo.yolo_config import YoloConfiguration
+from pl_bolts.models.detection.yolo.yolo_module import Yolo
+
+__all__ = ["YoloConfiguration", "Yolo"]

pl_bolts/models/detection/yolo/yolo_config.py

@@ -0,0 +1,265 @@
+import re
+from typing import List, Tuple
+from warnings import warn
+
+import torch.nn as nn
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+
+from pl_bolts.models.detection.yolo import yolo_layers
+
+
+class YoloConfiguration:
+    """
+    This class can be used to parse the configuration files of the Darknet YOLOv4 implementation.
+    The `get_network()` method returns a PyTorch module list that can be used to construct a YOLO
+    model.
+    """
+
+    def __init__(self, path: str):
+        """
+        Saves the variables from the first configuration section to attributes of this object, and
+        the rest of the sections to the `layer_configs` list.
+
+        Args:
+            path: Path to a configuration file
+        """
+        with open(path, 'r') as config_file:
+            sections = self._read_file(config_file)
+
+        if len(sections) < 2:
+            raise MisconfigurationException("The model configuration file should include at least two sections.")
+
+        self.__dict__.update(sections[0])
+        self.global_config = sections[0]
+        self.layer_configs = sections[1:]
+
+    def get_network(self) -> nn.ModuleList:
+        """
+        Iterates through the layers from the configuration and creates corresponding PyTorch
+        modules. Returns the network structure that can be used to create a YOLO model.
+
+        Returns:
+            modules: A `nn.ModuleList` that defines the YOLO network.
+        """
+        result = nn.ModuleList()
+        num_inputs = [3]  # Number of channels in the input of every layer up to the current layer
+        for layer_config in self.layer_configs:
+            config = {**self.global_config, **layer_config}
+            module, num_outputs = _create_layer(config, num_inputs)
+            result.append(module)
+            num_inputs.append(num_outputs)
+        return result
+
+    def _read_file(self, config_file):
+        """
+        Reads a YOLOv4 network configuration file and returns a list of configuration sections.
+
+        Args:
+            config_file (iterable over lines): The configuration file to read.
+
+        Returns:
+            sections (List[dict]): A list of configuration sections.
+        """
+        section_re = re.compile(r'\[([^]]+)\]')
+        list_variables = ('layers', 'anchors', 'mask', 'scales')
+        variable_types = {
+            'activation': str,
+            'anchors': int,
+            'angle': float,
+            'batch': int,
+            'batch_normalize': bool,
+            'beta_nms': float,
+            'burn_in': int,
+            'channels': int,
+            'classes': int,
+            'cls_normalizer': float,
+            'decay': float,
+            'exposure': float,
+            'filters': int,
+            'from': int,
+            'groups': int,
+            'group_id': int,
+            'height': int,
+            'hue': float,
+            'ignore_thresh': float,
+            'iou_loss': str,
+            'iou_normalizer': float,
+            'iou_thresh': float,
+            'jitter': float,
+            'layers': int,
+            'learning_rate': float,
+            'mask': int,
+            'max_batches': int,
+            'max_delta': float,
+            'momentum': float,
+            'mosaic': bool,
+            'nms_kind': str,
+            'num': int,
+            'obj_normalizer': float,
+            'pad': bool,
+            'policy': str,
+            'random': bool,
+            'resize': float,
+            'saturation': float,
+            'scales': float,
+            'scale_x_y': float,
+            'size': int,
+            'steps': str,
+            'stride': int,
+            'subdivisions': int,
+            'truth_thresh': float,
+            'width': int
+        }
+
+        section = None
+        sections = []
+
+        def convert(key, value):
+            """Converts a value to the correct type based on key."""
+            if key not in variable_types:
+                warn('Unknown YOLO configuration variable: ' + key)
+                return key, value
+            if key in list_variables:
+                value = [variable_types[key](v) for v in value.split(',')]
+            else:
+                value = variable_types[key](value)
+            return key, value
+
+        for line in config_file:
+            line = line.strip()
+            if (not line) or (line[0] == '#'):
+                continue
+
+            section_match = section_re.match(line)
+            if section_match:
+                if section is not None:
+                    sections.append(section)
+                section = {'type': section_match.group(1)}
+            else:
+                key, value = line.split('=')
+                key = key.rstrip()
+                value = value.lstrip()
+                key, value = convert(key, value)
+                section[key] = value
+        if section is not None:
+            sections.append(section)
+
+        return sections
+
+
+def _create_layer(config: dict, num_inputs: List[int]) -> Tuple[nn.Module, int]:
+    """
+    Calls one of the `_create_<layertype>(config, num_inputs)` functions to create a PyTorch
+    module from the layer config.
+
+    Args:
+        config: Dictionary of configuration options for this layer.
+        num_inputs: Number of channels in the input of every layer up to this layer.
+
+    Returns:
+        module (:class:`~torch.nn.Module`), num_outputs (int): The created PyTorch module and the
+            number of channels in its output.
+    """
+    create_func = {
+        'convolutional': _create_convolutional,
+        'maxpool': _create_maxpool,
+        'route': _create_route,
+        'shortcut': _create_shortcut,
+        'upsample': _create_upsample,
+        'yolo': _create_yolo
+    }
+    return create_func[config['type']](config, num_inputs)
+
+
+def _create_convolutional(config, num_inputs):
+    module = nn.Sequential()
+
+    batch_normalize = config.get('batch_normalize', False)
+    padding = (config['size'] - 1) // 2 if config['pad'] else 0
+
+    conv = nn.Conv2d(
+        num_inputs[-1], config['filters'], config['size'], config['stride'], padding, bias=not batch_normalize
+    )
+    module.add_module('conv', conv)
+
+    if batch_normalize:
+        bn = nn.BatchNorm2d(config['filters'])
+        module.add_module('bn', bn)
+
+    if config['activation'] == 'leaky':
+        leakyrelu = nn.LeakyReLU(0.1, inplace=True)
+        module.add_module('leakyrelu', leakyrelu)
+    elif config['activation'] == 'mish':
+        mish = yolo_layers.Mish()
+        module.add_module('mish', mish)
+
+    return module, config['filters']
+
+
+def _create_maxpool(config, num_inputs):
+    padding = (config['size'] - 1) // 2
+    module = nn.MaxPool2d(config['size'], config['stride'], padding)
+    return module, num_inputs[-1]
+
+
+def _create_route(config, num_inputs):
+    num_chunks = config.get('groups', 1)
+    chunk_idx = config.get('group_id', 0)
+
+    # 0 is the first layer, -1 is the previous layer
+    last = len(num_inputs) - 1
+    source_layers = [layer if layer >= 0 else last + layer for layer in config['layers']]
+
+    module = yolo_layers.RouteLayer(source_layers, num_chunks, chunk_idx)
+
+    # The number of outputs of a source layer is the number of inputs of the next layer.
+    num_outputs = sum(num_inputs[layer + 1] // num_chunks for layer in source_layers)
+
+    return module, num_outputs
+
+
+def _create_shortcut(config, num_inputs):
+    module = yolo_layers.ShortcutLayer(config['from'])
+    return module, num_inputs[-1]
+
+
+def _create_upsample(config, num_inputs):
+    module = nn.Upsample(scale_factor=config["stride"], mode='nearest')
+    return module, num_inputs[-1]
+
+
+def _create_yolo(config, num_inputs):
+    # The "anchors" list alternates width and height.
+    anchor_dims = config['anchors']
+    anchor_dims = [(anchor_dims[i], anchor_dims[i + 1]) for i in range(0, len(anchor_dims), 2)]
+
+    xy_scale = config.get('scale_x_y', 1.0)
+    ignore_threshold = config.get('ignore_thresh', 1.0)
+    overlap_loss_multiplier = config.get('iou_normalizer', 1.0)
+    class_loss_multiplier = config.get('cls_normalizer', 1.0)
+    confidence_loss_multiplier = config.get('obj_normalizer', 1.0)
+
+    overlap_loss_name = config.get('iou_loss', 'mse')
+    if overlap_loss_name == 'mse':
+        overlap_loss_func = nn.MSELoss(reduction='none')
+    elif overlap_loss_name == 'giou':
+        overlap_loss_func = yolo_layers.GIoULoss()
+    else:
+        overlap_loss_func = yolo_layers.IoULoss()
+
+    module = yolo_layers.DetectionLayer(
+        num_classes=config['classes'],
+        image_width=config['width'],
+        image_height=config['height'],
+        anchor_dims=anchor_dims,
+        anchor_ids=config['mask'],
+        xy_scale=xy_scale,
+        ignore_threshold=ignore_threshold,
+        overlap_loss_func=overlap_loss_func,
+        image_space_loss=overlap_loss_name != 'mse',
+        overlap_loss_multiplier=overlap_loss_multiplier,
+        class_loss_multiplier=class_loss_multiplier,
+        confidence_loss_multiplier=confidence_loss_multiplier
+    )
+
+    return module, num_inputs[-1]