Prep a set of ResNetV2 models with GroupNorm, EvoNormB0, EvoNormS0 for BN free model experiments on TPU and IPU

timm/models/layers/evo_norm.py

@@ -55,7 +55,7 @@ class EvoNormBatch2d(nn.Module):
 
 
 class EvoNormSample2d(nn.Module):
-    def __init__(self, num_features, apply_act=True, groups=8, eps=1e-5, drop_block=None):
+    def __init__(self, num_features, apply_act=True, groups=32, eps=1e-5, drop_block=None):
         super(EvoNormSample2d, self).__init__()
         self.apply_act = apply_act  # apply activation (non-linearity)
         self.groups = groups

timm/models/layers/norm_act.py

@@ -68,7 +68,7 @@ class BatchNormAct2d(nn.BatchNorm2d):
 
 class GroupNormAct(nn.GroupNorm):
     # NOTE num_channel and num_groups order flipped for easier layer swaps / binding of fixed args
-    def __init__(self, num_channels, num_groups, eps=1e-5, affine=True,
+    def __init__(self, num_channels, num_groups=32, eps=1e-5, affine=True,
                  apply_act=True, act_layer=nn.ReLU, inplace=True, drop_block=None):
         super(GroupNormAct, self).__init__(num_groups, num_channels, eps=eps, affine=affine)
         if isinstance(act_layer, str):

timm/models/resnetv2.py

@@ -120,6 +120,13 @@ default_cfgs = {
         interpolation='bicubic'),
     'resnetv2_152d': _cfg(
         interpolation='bicubic', first_conv='stem.conv1'),
+
+    'resnetv2_50d_gn': _cfg(
+        interpolation='bicubic', first_conv='stem.conv1'),
+    'resnetv2_50d_evob': _cfg(
+        interpolation='bicubic', first_conv='stem.conv1'),
+    'resnetv2_50d_evos': _cfg(
+        interpolation='bicubic', first_conv='stem.conv1'),
 }
 
 
@@ -639,19 +646,27 @@ def resnetv2_152d(pretrained=False, **kwargs):
         stem_type='deep', avg_down=True, **kwargs)
 
 
-# @register_model
-# def resnetv2_50ebd(pretrained=False, **kwargs):
-#     # FIXME for testing w/ TPU + PyTorch XLA
-#     return _create_resnetv2(
-#         'resnetv2_50d', pretrained=pretrained,
-#         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=EvoNormBatch2d,
-#         stem_type='deep', avg_down=True, **kwargs)
-#
-#
-# @register_model
-# def resnetv2_50esd(pretrained=False, **kwargs):
-#     # FIXME for testing w/ TPU + PyTorch XLA
-#     return _create_resnetv2(
-#         'resnetv2_50d', pretrained=pretrained,
-#         layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=EvoNormSample2d,
-#         stem_type='deep', avg_down=True, **kwargs)
+# Experimental configs (may change / be removed)
+
+@register_model
+def resnetv2_50d_gn(pretrained=False, **kwargs):
+    return _create_resnetv2(
+        'resnetv2_50d_gn', pretrained=pretrained,
+        layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=GroupNormAct,
+        stem_type='deep', avg_down=True, **kwargs)
+
+
+@register_model
+def resnetv2_50d_evob(pretrained=False, **kwargs):
+    return _create_resnetv2(
+        'resnetv2_50d_evob', pretrained=pretrained,
+        layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=EvoNormBatch2d,
+        stem_type='deep', avg_down=True, **kwargs)
+
+
+@register_model
+def resnetv2_50d_evos(pretrained=False, **kwargs):
+    return _create_resnetv2(
+        'resnetv2_50d_evos', pretrained=pretrained,
+        layers=[3, 4, 6, 3], conv_layer=create_conv2d, norm_layer=EvoNormSample2d,
+        stem_type='deep', avg_down=True, **kwargs)