Adding some configs to sknet, incl ResNet50 variants from 'Compounding ... Assembled Techniques' paper and original SKNet50

timm/models/sknet.py

@@ -22,7 +22,10 @@ def _cfg(url='', **kwargs):
 
 default_cfgs = {
     'skresnet18': _cfg(url=''),
-    'skresnet26d': _cfg()
+    'skresnet26d': _cfg(),
+    'skresnet50': _cfg(),
+    'skresnet50d': _cfg(),
+    'skresnext50_32x4d': _cfg(),
 }
 
 
@@ -131,6 +134,41 @@ class SelectiveKernelBottleneck(nn.Module):
         return x
 
 
+@register_model
+def skresnet18(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
+    """Constructs a ResNet-18 model.
+    """
+    default_cfg = default_cfgs['skresnet18']
+    sk_kwargs = dict(
+        min_attn_channels=16,
+    )
+    model = ResNet(
+        SelectiveKernelBasic, [2, 2, 2, 2], num_classes=num_classes, in_chans=in_chans,
+        block_args=dict(sk_kwargs=sk_kwargs), **kwargs)
+    model.default_cfg = default_cfg
+    if pretrained:
+        load_pretrained(model, default_cfg, num_classes, in_chans)
+    return model
+
+
+@register_model
+def sksresnet18(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
+    """Constructs a ResNet-18 model.
+    """
+    default_cfg = default_cfgs['skresnet18']
+    sk_kwargs = dict(
+        min_attn_channels=16,
+        split_input=True
+    )
+    model = ResNet(
+        SelectiveKernelBasic, [2, 2, 2, 2], num_classes=num_classes, in_chans=in_chans,
+        block_args=dict(sk_kwargs=sk_kwargs), **kwargs)
+    model.default_cfg = default_cfg
+    if pretrained:
+        load_pretrained(model, default_cfg, num_classes, in_chans)
+    return model
+
+
 @register_model
 def skresnet26d(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
     """Constructs a ResNet-26 model.
@@ -150,15 +188,17 @@ def skresnet26d(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
 
 
 @register_model
-def skresnet18(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
+def skresnet50(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
-    """Constructs a ResNet-18 model.
+    """Constructs a Select Kernel ResNet-50 model.
+    Based on config in "Compounding the Performance Improvements of Assembled Techniques in a
+    Convolutional Neural Network"
     """
-    default_cfg = default_cfgs['skresnet18']
     sk_kwargs = dict(
-        min_attn_channels=16,
+        attn_reduction=2,
     )
+    default_cfg = default_cfgs['skresnet50']
     model = ResNet(
-        SelectiveKernelBasic, [2, 2, 2, 2], num_classes=num_classes, in_chans=in_chans,
+        SelectiveKernelBottleneck, [3, 4, 6, 3], num_classes=num_classes, in_chans=in_chans,
         block_args=dict(sk_kwargs=sk_kwargs), **kwargs)
     model.default_cfg = default_cfg
     if pretrained:
@@ -167,17 +207,33 @@ def skresnet18(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
 
 
 @register_model
-def sksresnet18(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
+def skresnet50d(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
-    """Constructs a ResNet-18 model.
+    """Constructs a Select Kernel ResNet-50-D model.
+    Based on config in "Compounding the Performance Improvements of Assembled Techniques in a
+    Convolutional Neural Network"
     """
-    default_cfg = default_cfgs['skresnet18']
     sk_kwargs = dict(
-        min_attn_channels=16,
+        attn_reduction=2,
-        split_input=True
     )
+    default_cfg = default_cfgs['skresnet50d']
     model = ResNet(
-        SelectiveKernelBasic, [2, 2, 2, 2], num_classes=num_classes, in_chans=in_chans,
+        SelectiveKernelBottleneck, [3, 4, 6, 3], stem_width=32, stem_type='deep', avg_down=True,
-        block_args=dict(sk_kwargs=sk_kwargs), **kwargs)
+        num_classes=num_classes, in_chans=in_chans, block_args=dict(sk_kwargs=sk_kwargs), **kwargs)
+    model.default_cfg = default_cfg
+    if pretrained:
+        load_pretrained(model, default_cfg, num_classes, in_chans)
+    return model
+
+
+@register_model
+def skresnext50_32x4d(pretrained=False, num_classes=1000, in_chans=3, **kwargs):
+    """Constructs a Select Kernel ResNeXt50-32x4d model. This should be equivalent to
+    the SKNet50 model in the Select Kernel Paper
+    """
+    default_cfg = default_cfgs['skresnext50_32x4d']
+    model = ResNet(
+        SelectiveKernelBottleneck, [3, 4, 6, 3], cardinality=32, base_width=4,
+        num_classes=num_classes, in_chans=in_chans, **kwargs)
     model.default_cfg = default_cfg
     if pretrained:
         load_pretrained(model, default_cfg, num_classes, in_chans)