feat: Adds support of YOLOv4 detection model

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
 
 
 

docs/source/models.rst

@@ -112,6 +112,8 @@ YOLO
 
 .. autofunction:: yolov2
 
+.. autofunction:: yolov4
+
 
 Semantic Segmentation
 =====================

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

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