Repository: airctic/icevision

Issue:
Add ResNeSt backbones
## 🚀 Feature
ResNeSt aims to replace a ResNet backbones in both classification and object detection tasks and substantially improve accuracy.
It would be beneficial for IceVision to support the different ResNeSt architectures.

Paper: ResNeSt: Split-Attention Networks - https://arxiv.org/pdf/2004.08955.pdf 

Discord discussion: https://discordapp.com/channels/735877944085446747/735877944517591151/780303631738994700


Discussion:

Resolving patch:
diff --git a/icevision/backbones/__init__.py b/icevision/backbones/__init__.py
--- a/icevision/backbones/__init__.py
+++ b/icevision/backbones/__init__.py
@@ -3,3 +3,4 @@
 from icevision.backbones.mobilenet import *
 
 import icevision.backbones.resnet_fpn
+import icevision.backbones.resnest_fpn
diff --git a/icevision/backbones/resnest_fpn.py b/icevision/backbones/resnest_fpn.py
new file mode 100644
--- /dev/null
+++ b/icevision/backbones/resnest_fpn.py
@@ -0,0 +1,120 @@
+__all__ = [
+    "resnest50",
+    "resnest101",
+    "resnest200",
+    "resnest269",
+]
+
+from icevision.imports import *
+from icevision.utils import *
+from icevision.backbones.resnet_fpn import patch_param_groups
+from torchvision.ops.feature_pyramid_network import LastLevelMaxPool
+from torchvision.models.detection.backbone_utils import (
+    resnet_fpn_backbone,
+    BackboneWithFPN,
+)
+import resnest.torch
+
+
+# ResNeSt Backbones
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+
+
+def resnest_fpn_backbone(
+    resnest_name,
+    pretrained,
+    # norm_layer=misc_nn_ops.FrozenBatchNorm2d,
+    trainable_layers=3,
+    returned_layers=None,
+    extra_blocks=None,
+):
+    """
+    Constructs a specified ResNet backbone with FPN on top. Freezes the specified number of layers in the backbone.
+
+    Examples::
+
+        >>> from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
+        >>> backbone = resnet_fpn_backbone('resnet50', pretrained=True, trainable_layers=3)
+        >>> # get some dummy image
+        >>> x = torch.rand(1,3,64,64)
+        >>> # compute the output
+        >>> output = backbone(x)
+        >>> print([(k, v.shape) for k, v in output.items()])
+        >>> # returns
+        >>>   [('0', torch.Size([1, 256, 16, 16])),
+        >>>    ('1', torch.Size([1, 256, 8, 8])),
+        >>>    ('2', torch.Size([1, 256, 4, 4])),
+        >>>    ('3', torch.Size([1, 256, 2, 2])),
+        >>>    ('pool', torch.Size([1, 256, 1, 1]))]
+
+    Arguments:
+        backbone_name (string): resnet architecture. Possible values are 'ResNet', 'resnet18', 'resnet34', 'resnet50',
+             'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', 'wide_resnet101_2'
+        norm_layer (torchvision.ops): it is recommended to use the default value. For details visit:
+            (https://github.com/facebookresearch/maskrcnn-benchmark/issues/267)
+        pretrained (bool): If True, returns a model with backbone pre-trained on Imagenet
+        trainable_layers (int): number of trainable (not frozen) resnet layers starting from final block.
+            Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
+    """
+    backbone = resnest_name(pretrained=pretrained)
+    # backbone = resnest50(pretrained=pretrained)
+
+    backbone.fc = Identity()
+    backbone.avgpool = Identity()
+
+    # select layers that wont be frozen
+    assert trainable_layers <= 5 and trainable_layers >= 0
+    layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][
+        :trainable_layers
+    ]
+    # freeze layers only if pretrained backbone is used
+    for name, parameter in backbone.named_parameters():
+        if all([not name.startswith(layer) for layer in layers_to_train]):
+            parameter.requires_grad_(False)
+
+    if extra_blocks is None:
+        extra_blocks = LastLevelMaxPool()
+
+    if returned_layers is None:
+        returned_layers = [1, 2, 3, 4]
+    assert min(returned_layers) > 0 and max(returned_layers) < 5
+    return_layers = {f"layer{k}": str(v) for v, k in enumerate(returned_layers)}
+
+    in_channels_stage2 = backbone.inplanes // 8
+    in_channels_list = [in_channels_stage2 * 2 ** (i - 1) for i in returned_layers]
+    out_channels = 256
+    return BackboneWithFPN(
+        backbone,
+        return_layers,
+        in_channels_list,
+        out_channels,
+        extra_blocks=extra_blocks,
+    )
+
+
+def _resnest_fpn(name, pretrained: bool = True, **kwargs):
+    model = resnest_fpn_backbone(resnest_name=name, pretrained=pretrained, **kwargs)
+    patch_param_groups(model)
+
+    return model
+
+
+def resnest50(pretrained: bool = True, **kwargs):
+    return _resnest_fpn(resnest.torch.resnest50, pretrained=pretrained, **kwargs)
+
+
+def resnest101(pretrained: bool = True, **kwargs):
+    return _resnest_fpn(resnest.torch.resnest101, pretrained=pretrained, **kwargs)
+
+
+def resnest200(pretrained: bool = True, **kwargs):
+    return _resnest_fpn(resnest.torch.resnest200, pretrained=pretrained, **kwargs)
+
+
+def resnest269(pretrained: bool = True, **kwargs):
+    return _resnest_fpn(resnest.torch.resnest269, pretrained=pretrained, **kwargs)


Test patch: