feat: Adds support of YOLOv4 detection model (#78)

* feat: Updated Darknet URLs * docs: Updated Readmes * fix: Fixed backbone norm layer * feat: Added SPP & PAN layers * feat: Added YOLOv4 neck * style: Fixed lint * feat: Added Self Attention Module * docs: Fixed docstring * style: Added extra_repr to GlobalAvgPool2d * feat: Added implementation of YOLOv4 * refactor: Updated addcdiv usage to reflect torch changes * refactor: Switched anchors to buffer * style: Fixed lint * feat: Added compatibility for list of Tensor as input * fix: Passed down gt_boxes and gt_labels * style: Fixed lint * refactor: Reshaped model outputs * feat: Added implementation of Generalized IoU * test: Updated unittests * docs: Updated documentation * fix: Fixed YOLOv4 loss * test: Updated unittests for models * docs: Updated documentation * test: Added check to make sure loss can be backpropagated * fix: Fixed YOLO loss * docs: Updated docstring * feat: Updated default backbone parameters * refactor: Refactored loss computation * refactor: Updated global avg poolings in models * style: Fixed lint * fix: Fixed missed import * chore: Fixed imports * fix: Fixed typo * refactor: Refactored SPP * fix: Fixed anchor selection * refactor: Switched bbox loss to CIoU * refactor: Moved YOLOv4 to yolo submodule * docs: Updated readme * fix: Fixed YOLO initialization * style: Fixed codacy lint * fix: Fixed imports * fix: Fixed YOLO head * feat: Added zero initialization of last conv * refactor: Refactored anchor * feat: Added default initialization * fix: Fixed coordinate scaling * fix: Fixed anchor registration * fix: Fixed YOLOv4 loss * style: Fixed lint * style: Fixed codacy syntax * refactor: Refactored yolo layer
frgfm · Sep 2, 2020 · 5e54efa · 5e54efa
1 parent e133e05
commit 5e54efa
Show file tree

Hide file tree

Showing 12 changed files with 563 additions and 49 deletions.
diff --git a/README.md b/README.md
@@ -66,7 +66,7 @@ conda install -c frgfm pylocron
 ##### Main features
 
 - Classification: [Res2Net](https://arxiv.org/abs/1904.01169) (based on the great [implementation](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/res2net.py) from Ross Wightman), [Darknet-24](https://pjreddie.com/media/files/papers/yolo_1.pdf), [Darknet-19](https://pjreddie.com/media/files/papers/YOLO9000.pdf), [Darknet-53](https://pjreddie.com/media/files/papers/YOLOv3.pdf), [CSPDarknet-53](<https://arxiv.org/abs/1911.11929>), [ResNet](https://arxiv.org/abs/1512.03385), [ResNeXt](https://arxiv.org/abs/1611.05431), [TridentNet](https://arxiv.org/abs/1901.01892), [PyConvResNet](https://arxiv.org/abs/2006.11538), [ReXNet](https://arxiv.org/abs/2007.00992).
-- Detection: [YOLOv1](https://pjreddie.com/media/files/papers/yolo_1.pdf), [YOLOv2](https://pjreddie.com/media/files/papers/YOLO9000.pdf)
+- Detection: [YOLOv1](https://pjreddie.com/media/files/papers/yolo_1.pdf), [YOLOv2](https://pjreddie.com/media/files/papers/YOLO9000.pdf), [YOLOv4](https://arxiv.org/abs/2004.10934)
 - Segmentation: [U-Net](https://arxiv.org/abs/1505.04597), [UNet++](https://arxiv.org/abs/1807.10165), [UNet3+](https://arxiv.org/abs/2004.08790)
 
 ### ops
@@ -89,9 +89,9 @@ conda install -c frgfm pylocron
 
 The project is currently under development, here are the objectives for the next releases:
 
-- [ ] Standardize models: standardize models by task.
+- [x] Standardize models: standardize models by task.
+- [x] Reference scripts: add reference training scripts
 - [ ] Speed benchmark: compare `holocron.nn` functions execution speed.
-- [ ] Reference scripts: add reference training scripts
 
 
 

diff --git a/docs/source/models.rst b/docs/source/models.rst
@@ -112,6 +112,8 @@ YOLO
 
 .. autofunction:: yolov2
 
+.. autofunction:: yolov4
+
 
 Semantic Segmentation
 =====================

diff --git a/docs/source/ops.rst b/docs/source/ops.rst
@@ -11,5 +11,6 @@ holocron.ops
 Boxes
 -----
 
-.. autofunction:: box_diou
-.. autofunction:: box_ciou
+.. autofunction:: box_giou
+.. autofunction:: diou_loss
+.. autofunction:: ciou_loss
diff --git a/holocron/models/darknet.py b/holocron/models/darknet.py
@@ -14,7 +14,7 @@
 from ..nn.init import init_module
 from .utils import conv_sequence
 from .resnet import _ResBlock
-from holocron.nn import Mish, DropBlock2d
+from holocron.nn import Mish, DropBlock2d, GlobalAvgPool2d
 
 
 __all__ = ['DarknetV1', 'DarknetV2', 'DarknetV3', 'DarknetV4', 'darknet24', 'darknet19', 'darknet53', 'cspdarknet53']
@@ -23,13 +23,13 @@
 default_cfgs = {
     'darknet24': {'arch': 'DarknetV1',
                   'layout': [[192], [128, 256, 256, 512], [*([256, 512] * 4), 512, 1024], [512, 1024] * 2],
-                  'url': None},
+                  'url': 'https://github.com/frgfm/Holocron/releases/download/v0.1.2/darknet24_224-55729a5c.pth'},
     'darknet19': {'arch': 'DarknetV2',
                   'layout': [(64, 0), (128, 1), (256, 1), (512, 2), (1024, 2)],
-                  'url': None},
+                  'url': 'https://github.com/frgfm/Holocron/releases/download/v0.1.2/darknet19_224-a48304cd.pth'},
     'darknet53': {'arch': 'DarknetV3',
                   'layout': [(64, 1), (128, 2), (256, 8), (512, 8), (1024, 4)],
-                  'url': None},
+                  'url': 'https://github.com/frgfm/Holocron/releases/download/v0.1.2/darknet53_256-f57b8429.pth'},
     'cspdarknet53': {'arch': 'DarknetV4',
                      'layout': [(64, 1), (128, 2), (256, 8), (512, 8), (1024, 4)],
                      'url': None},
@@ -73,7 +73,7 @@ def __init__(self, layout, num_classes=10, in_channels=3, stem_channels=64,
         super().__init__(OrderedDict([
             ('features', DarknetBodyV1(layout, in_channels, stem_channels,
                                        act_layer, norm_layer, drop_layer, conv_layer)),
-            ('global_pool', nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten())),
+            ('pool', GlobalAvgPool2d(flatten=True)),
             ('classifier', nn.Linear(layout[2][-1], num_classes))]))
 
         init_module(self, 'leaky_relu')
@@ -139,7 +139,7 @@ def __init__(self, layout, num_classes=10, in_channels=3, stem_channels=32,
             ('features', DarknetBodyV2(layout, in_channels, stem_channels, False,
                                        act_layer, norm_layer, drop_layer, conv_layer)),
             ('classifier', nn.Conv2d(layout[-1][0], num_classes, 1)),
-            ('global_pool', nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten()))]))
+            ('pool', GlobalAvgPool2d(flatten=True))]))
 
         init_module(self, 'leaky_relu')
 
@@ -208,7 +208,7 @@ def __init__(self, layout, num_classes=10, in_channels=3, stem_channels=32,
         super().__init__(OrderedDict([
             ('features', DarknetBodyV3(layout, in_channels, stem_channels,
                                        act_layer, norm_layer, drop_layer, conv_layer)),
-            ('global_pool', nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten())),
+            ('pool', GlobalAvgPool2d(flatten=True)),
             ('classifier', nn.Linear(layout[-1][0], num_classes))]))
 
         init_module(self, 'leaky_relu')
@@ -297,7 +297,7 @@ def __init__(self, layout, num_classes=10, in_channels=3, stem_channels=32, num_
         super().__init__(OrderedDict([
             ('features', DarknetBodyV4(layout, in_channels, stem_channels, num_features,
                                        act_layer, norm_layer, drop_layer, conv_layer)),
-            ('global_pool', nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten())),
+            ('pool', GlobalAvgPool2d(flatten=True)),
             ('classifier', nn.Linear(layout[-1][0], num_classes))]))
 
         init_module(self, 'leaky_relu')