feat: Added implementation of Pyramidal convolution (#66)

* feat: Added pyramidal convolution * test: Updated unittests * docs: Updated documentation * docs: Updated readme * feat: Updated PyConv2d for more flexibility * test: Updated unittests * refactor: Refactored PyConv2d to handle expansion=1 * test: Updated unittests * refactor: Fixed kwargs passing to conv_sequence * fix: Fixed conv type passing in ResNet * feat: Added option to turn off stem pooling * style: Renamed PyConv attribute * fix: Fixed conv type passing to downsampling blocks * refactor: Changed block arguments passing * feat: Added PyConvResNet implementation * fix: Fixed Res2Net block argument passing * docs: Updated documentation * docs: Updated readme * test: Add unittest for pyconvresnet50 * style: Fixed lint * refactor: Changed block argument passing mechanism * feat: Added flexibility for grouping in PyConv * feat: Added PyConvHGResNet * test: Updating unittests * docs: Updated documentation * style: Fixed lint * fix: Fixed edge case of PyConv2d
frgfm · Jul 24, 2020 · f2e1e8c · f2e1e8c
1 parent 2be2531
commit f2e1e8c
Show file tree

Hide file tree

Showing 10 changed files with 227 additions and 42 deletions.
diff --git a/README.md b/README.md
@@ -58,14 +58,14 @@ conda install -c frgfm pylocron
 
 - Activation: [SiLU/Swish](https://arxiv.org/abs/1606.08415), [Mish](https://arxiv.org/abs/1908.08681), [HardMish](https://github.com/digantamisra98/H-Mish), [NLReLU](https://arxiv.org/abs/1908.03682)
 - Loss: [Focal Loss](https://arxiv.org/abs/1708.02002), MultiLabelCrossEntropy, [LabelSmoothingCrossEntropy](https://arxiv.org/pdf/1706.03762.pdf), [MixupLoss](https://arxiv.org/pdf/1710.09412.pdf)
-- Convolutions: [NormConv2d](https://arxiv.org/pdf/2005.05274v2.pdf), [Add2d](https://arxiv.org/pdf/1912.13200.pdf), [SlimConv2d](https://arxiv.org/pdf/2003.07469.pdf)
+- Convolutions: [NormConv2d](https://arxiv.org/pdf/2005.05274v2.pdf), [Add2d](https://arxiv.org/pdf/1912.13200.pdf), [SlimConv2d](https://arxiv.org/pdf/2003.07469.pdf), [PyConv2d](https://arxiv.org/abs/2006.11538)
 - Regularization: [DropBlock](https://arxiv.org/abs/1810.12890)
 
 ### models
 
 ##### 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), [darknet24](https://pjreddie.com/media/files/papers/yolo_1.pdf), [darknet19](https://pjreddie.com/media/files/papers/YOLO9000.pdf), [darknet53](https://pjreddie.com/media/files/papers/YOLOv3.pdf), [ResNet](https://arxiv.org/abs/1512.03385), [ResNeXt](https://arxiv.org/abs/1611.05431), [TridentNet](https://arxiv.org/abs/1901.01892), [ReXNet](https://arxiv.org/abs/2007.00992).
+- 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), [darknet24](https://pjreddie.com/media/files/papers/yolo_1.pdf), [darknet19](https://pjreddie.com/media/files/papers/YOLO9000.pdf), [darknet53](https://pjreddie.com/media/files/papers/YOLOv3.pdf), [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)
 - Segmentation: [U-Net](https://arxiv.org/abs/1505.04597), [UNet++](https://arxiv.org/abs/1807.10165), [UNet3+](https://arxiv.org/abs/2004.08790)
 

diff --git a/docs/source/models.rst b/docs/source/models.rst
@@ -58,6 +58,14 @@ TridentNet
 .. autofunction:: tridentnet50
 
 
+PyConvResNet
+------------
+
+.. autofunction:: pyconv_resnet50
+
+.. autofunction:: pyconvhg_resnet50
+
+
 ReXNet
 -------
 

diff --git a/docs/source/nn.rst b/docs/source/nn.rst
@@ -40,6 +40,8 @@ Convolution layers
 
 .. autoclass:: SlimConv2d
 
+.. autoclass:: PyConv2d
+
 Regularization layers
 ---------------------
 

diff --git a/holocron/models/__init__.py b/holocron/models/__init__.py
@@ -1,6 +1,7 @@
 from .resnet import *
 from .res2net import *
 from .tridentnet import *
+from .pyconvresnet import *
 from .rexnet import *
 from .darknet import *
 from .detection import *

diff --git a/holocron/models/pyconvresnet.py b/holocron/models/pyconvresnet.py
@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+
+"""
+Implementation of PyConvResNet
+"""
+
+import sys
+import logging
+import torch.nn as nn
+from torchvision.models.utils import load_state_dict_from_url
+from holocron.nn import PyConv2d
+from .resnet import ResNet, _ResBlock
+from .utils import conv_sequence
+
+
+__all__ = ['PyBottleneck', 'pyconv_resnet50', 'pyconvhg_resnet50']
+
+
+default_cfgs = {
+    'pyconv_resnet50': {'block': 'PyBottleneck', 'num_blocks': [3, 4, 6, 3], 'out_chans': [64, 128, 256, 512],
+                        'width_per_group': 64,
+                        'groups': [[1, 4, 8, 16], [1, 4, 8], [1, 4], [1]],
+                        'url': None},
+    'pyconvhg_resnet50': {'block': 'PyHGBottleneck', 'num_blocks': [3, 4, 6, 3], 'out_chans': [128, 256, 512, 1024],
+                          'width_per_group': 2,
+                          'groups': [[32, 32, 32, 32], [32, 64, 64], [32, 64], [32]],
+                          'url': None},
+}
+
+
+class PyBottleneck(_ResBlock):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=None, base_width=64, dilation=1,
+                 act_layer=None, norm_layer=None, drop_layer=None, num_levels=2, **kwargs):
+
+        width = int(planes * (base_width / 64.)) * min(groups)
+
+        super().__init__(
+            [*conv_sequence(inplanes, width, act_layer, norm_layer, drop_layer, kernel_size=1,
+                            stride=1, bias=False, **kwargs),
+             *conv_sequence(width, width, act_layer, norm_layer, drop_layer, conv_layer=PyConv2d, kernel_size=3,
+                            stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation,
+                            num_levels=num_levels, **kwargs),
+             *conv_sequence(width, planes * self.expansion, None, norm_layer, drop_layer, kernel_size=1,
+                            stride=1, bias=False, **kwargs)],
+            downsample, act_layer)
+
+
+class PyHGBottleneck(PyBottleneck):
+    expansion = 2
+
+
+def _pyconvresnet(arch, pretrained, progress, **kwargs):
+
+    # Retrieve the correct block type
+    block = sys.modules[__name__].__dict__[default_cfgs[arch]['block']]
+    # Build the model
+    model = ResNet(block, default_cfgs[arch]['num_blocks'], default_cfgs[arch]['out_chans'], stem_pool=False,
+                   width_per_group=default_cfgs[arch]['width_per_group'],
+                   block_args=[dict(num_levels=len(group), groups=group)
+                               for group in default_cfgs[arch]['groups']], **kwargs)
+    # Load pretrained parameters
+    if pretrained:
+        if default_cfgs[arch]['url'] is None:
+            logging.warning(f"Invalid model URL for {arch}, using default initialization.")
+        else:
+            state_dict = load_state_dict_from_url(default_cfgs[arch]['url'],
+                                                  progress=progress)
+            model.load_state_dict(state_dict)
+
+    return model
+
+
+def pyconv_resnet50(pretrained=False, progress=True, **kwargs):
+    """PyConvResNet-50 from `"Pyramidal Convolution: Rethinking Convolutional Neural Networks
+    for Visual Recognition" <https://arxiv.org/pdf/2006.11538.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+
+    Returns:
+        torch.nn.Module: classification model
+    """
+
+    return _pyconvresnet('pyconv_resnet50', pretrained, progress, **kwargs)
+
+
+def pyconvhg_resnet50(pretrained=False, progress=True, **kwargs):
+    """PyConvHGResNet-50 from `"Pyramidal Convolution: Rethinking Convolutional Neural Networks
+    for Visual Recognition" <https://arxiv.org/pdf/2006.11538.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+
+    Returns:
+        torch.nn.Module: classification model
+    """
+
+    return _pyconvresnet('pyconvhg_resnet50', pretrained, progress, **kwargs)
diff --git a/holocron/models/res2net.py b/holocron/models/res2net.py
@@ -94,7 +94,9 @@ def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
 def _res2net(arch, pretrained, progress, **kwargs):
     # Build the model
     model = ResNet(Bottle2neck, default_cfgs[arch]['num_blocks'], [64, 128, 256, 512],
-                   width_per_group=default_cfgs[arch]['width_per_group'], scale=default_cfgs[arch]['scale'], **kwargs)
+                   width_per_group=default_cfgs[arch]['width_per_group'],
+                   block_args=dict(scale=default_cfgs[arch]['scale']),
+                   **kwargs)
     # Load pretrained parameters
     if pretrained:
         if default_cfgs[arch]['url'] is None:

diff --git a/holocron/models/resnet.py b/holocron/models/resnet.py
@@ -69,31 +69,31 @@ class BasicBlock(_ResBlock):
 
     expansion = 1
 
-    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
-                 dilation=1, act_layer=None, norm_layer=None, drop_layer=None):
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1,
+                 act_layer=None, norm_layer=None, drop_layer=None, conv_layer=None, **kwargs):
         super().__init__(
-            [*conv_sequence(inplanes, planes, act_layer, norm_layer, drop_layer, kernel_size=3, stride=stride,
-                            padding=dilation, groups=groups, bias=False, dilation=dilation),
-             *conv_sequence(planes, planes, None, norm_layer, drop_layer, kernel_size=3, stride=1,
-                            padding=dilation, groups=groups, bias=False, dilation=dilation)],
+            [*conv_sequence(inplanes, planes, act_layer, norm_layer, drop_layer, conv_layer, kernel_size=3,
+                            stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation, **kwargs),
+             *conv_sequence(planes, planes, None, norm_layer, drop_layer, conv_layer, kernel_size=3,
+                            stride=1, padding=dilation, groups=groups, bias=False, dilation=dilation, **kwargs)],
             downsample, act_layer)
 
 
 class Bottleneck(_ResBlock):
 
     expansion = 4
 
-    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
-                 base_width=64, dilation=1, act_layer=None, norm_layer=None, drop_layer=None):
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1,
+                 act_layer=None, norm_layer=None, drop_layer=None, conv_layer=None, **kwargs):
 
         width = int(planes * (base_width / 64.)) * groups
         super().__init__(
-            [*conv_sequence(inplanes, width, act_layer, norm_layer, drop_layer, kernel_size=1,
-                            stride=1, bias=False),
-             *conv_sequence(width, width, act_layer, norm_layer, drop_layer, kernel_size=3, stride=stride,
-                            padding=dilation, groups=groups, bias=False, dilation=dilation),
-             *conv_sequence(width, planes * self.expansion, None, norm_layer, drop_layer, kernel_size=1,
-                            stride=1, bias=False)],
+            [*conv_sequence(inplanes, width, act_layer, norm_layer, drop_layer, conv_layer, kernel_size=1,
+                            stride=1, bias=False, **kwargs),
+             *conv_sequence(width, width, act_layer, norm_layer, drop_layer, conv_layer, kernel_size=3,
+                            stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation, **kwargs),
+             *conv_sequence(width, planes * self.expansion, None, norm_layer, drop_layer, conv_layer, kernel_size=1,
+                            stride=1, bias=False, **kwargs)],
             downsample, act_layer)
 
 
@@ -111,9 +111,9 @@ def forward(self, x):
 
 class ResNet(nn.Sequential):
     def __init__(self, block, num_blocks, planes, num_classes=10, in_channels=3, zero_init_residual=False,
-                 groups=1, width_per_group=64, conv_layer=None,
-                 act_layer=None, norm_layer=None, drop_layer=None, deep_stem=False, avg_downsample=False,
-                 num_repeats=1, **kwargs):
+                 width_per_group=64,
+                 conv_layer=None, act_layer=None, norm_layer=None, drop_layer=None, deep_stem=False, stem_pool=True,
+                 avg_downsample=False, num_repeats=1, block_args=None):
 
         if conv_layer is None:
             conv_layer = nn.Conv2d
@@ -126,27 +126,34 @@ def __init__(self, block, num_blocks, planes, num_classes=10, in_channels=3, zer
         in_planes = 64
         # Deep stem from ResNet-C
         if deep_stem:
-            _layers = [*conv_sequence(in_channels, in_planes // 2, act_layer, norm_layer, drop_layer,
+            _layers = [*conv_sequence(in_channels, in_planes // 2, act_layer, norm_layer, drop_layer, conv_layer,
                                       kernel_size=3, stride=2, padding=1, bias=False),
-                       *conv_sequence(in_planes // 2, in_planes // 2, act_layer, norm_layer, drop_layer,
+                       *conv_sequence(in_planes // 2, in_planes // 2, act_layer, norm_layer, drop_layer, conv_layer,
                                       kernel_size=3, stride=1, padding=1, bias=False),
-                       *conv_sequence(in_planes // 2, in_planes, act_layer, norm_layer, drop_layer,
+                       *conv_sequence(in_planes // 2, in_planes, act_layer, norm_layer, drop_layer, conv_layer,
                                       kernel_size=3, stride=1, padding=1, bias=False)]
         else:
-            _layers = conv_sequence(in_channels, in_planes, act_layer, norm_layer, drop_layer,
+            _layers = conv_sequence(in_channels, in_planes, act_layer, norm_layer, drop_layer, conv_layer,
                                     kernel_size=7, stride=2, padding=3, bias=False)
-        _layers.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
+        if stem_pool:
+            _layers.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
 
         # Optional tensor repetitions along channel axis (mainly for TridentNet)
         if num_repeats > 1:
             _layers.append(ChannelRepeat(num_repeats))
 
         # Consecutive convolutional blocks
         stride = 1
-        for _num_blocks, _planes in zip(num_blocks, planes):
-            _layers.append(self._make_layer(block, _num_blocks, in_planes, _planes, stride, groups, width_per_group,
+        # Block args
+        if block_args is None:
+            block_args = dict(groups=1)
+        if not isinstance(block_args, list):
+            block_args = [block_args] * len(num_blocks)
+        for _num_blocks, _planes, _block_args in zip(num_blocks, planes, block_args):
+            _layers.append(self._make_layer(block, _num_blocks, in_planes, _planes, stride, width_per_group,
                                             act_layer=act_layer, norm_layer=norm_layer, drop_layer=drop_layer,
-                                            avg_downsample=avg_downsample, num_repeats=num_repeats, **kwargs))
+                                            avg_downsample=avg_downsample, num_repeats=num_repeats,
+                                            block_args=_block_args))
             in_planes = block.expansion * _planes
             stride = 2
 
@@ -168,9 +175,9 @@ def __init__(self, block, num_blocks, planes, num_classes=10, in_channels=3, zer
                     m.convs[1][1].weight.data.zero_()
 
     @staticmethod
-    def _make_layer(block, num_blocks, in_planes, planes, stride=1, groups=1, width_per_group=64,
-                    norm_layer=None, act_layer=None, drop_layer=None, avg_downsample=False, num_repeats=1,
-                    **kwargs):
+    def _make_layer(block, num_blocks, in_planes, planes, stride=1, width_per_group=64,
+                    act_layer=None, norm_layer=None, drop_layer=None, conv_layer=None,
+                    avg_downsample=False, num_repeats=1, block_args=None):
 
         downsample = None
         if stride != 1 or in_planes != planes * block.expansion:
@@ -179,19 +186,21 @@ def _make_layer(block, num_blocks, in_planes, planes, stride=1, groups=1, width_
                 downsample = nn.Sequential(nn.AvgPool2d(stride, ceil_mode=True, count_include_pad=False),
                                            *conv_sequence(num_repeats * in_planes,
                                                           num_repeats * planes * block.expansion,
-                                                          None, norm_layer, drop_layer,
+                                                          None, norm_layer, drop_layer, conv_layer,
                                                           kernel_size=1, stride=1, bias=False))
             else:
                 downsample = nn.Sequential(*conv_sequence(num_repeats * in_planes,
                                                           num_repeats * planes * block.expansion,
-                                                          None, norm_layer, drop_layer,
+                                                          None, norm_layer, drop_layer, conv_layer,
                                                           kernel_size=1, stride=stride, bias=False))
-        layers = [block(in_planes, planes, stride, downsample, groups, width_per_group,
-                        act_layer=act_layer, norm_layer=norm_layer, drop_layer=drop_layer, **kwargs)]
+        if block_args is None:
+            block_args = {}
+        layers = [block(in_planes, planes, stride, downsample, base_width=width_per_group,
+                        act_layer=act_layer, norm_layer=norm_layer, drop_layer=drop_layer, **block_args)]
 
         for _ in range(num_blocks - 1):
-            layers.append(block(block.expansion * planes, planes, 1, None, groups, width_per_group,
-                                act_layer=act_layer, norm_layer=norm_layer, drop_layer=drop_layer, **kwargs))
+            layers.append(block(block.expansion * planes, planes, 1, None, base_width=width_per_group,
+                                act_layer=act_layer, norm_layer=norm_layer, drop_layer=drop_layer, **block_args))
 
         return nn.Sequential(*layers)
 
@@ -318,9 +327,9 @@ def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
         torch.nn.Module: classification model
     """
 
-    kwargs['groups'] = 32
     kwargs['width_per_group'] = 4
-    return _resnet('resnext50_32x4d', pretrained, progress, **kwargs)
+    block_args = dict(groups=32)
+    return _resnet('resnext50_32x4d', pretrained, progress, block_args=block_args, **kwargs)
 
 
 def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
@@ -335,6 +344,6 @@ def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
         torch.nn.Module: classification model
     """
 
-    kwargs['groups'] = 32
     kwargs['width_per_group'] = 8
-    return _resnet('resnext101_32x8d', pretrained, progress, **kwargs)
+    block_args = dict(groups=32)
+    return _resnet('resnext101_32x8d', pretrained, progress, block_args=block_args, **kwargs)