From af607b75cc7aa91eaece8a4478939e33a17cd1e9 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Fri, 19 Nov 2021 17:37:00 -0800 Subject: [PATCH] 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 | 2 +- timm/models/layers/norm_act.py | 2 +- timm/models/resnetv2.py | 47 ++++++++++++++++++++++------------ 3 files changed, 33 insertions(+), 18 deletions(-) diff --git a/timm/models/layers/evo_norm.py b/timm/models/layers/evo_norm.py index ecc5fb61..50367f9b 100644 --- a/timm/models/layers/evo_norm.py +++ b/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 diff --git a/timm/models/layers/norm_act.py b/timm/models/layers/norm_act.py index 02cabe88..2e15181f 100644 --- a/timm/models/layers/norm_act.py +++ b/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): diff --git a/timm/models/resnetv2.py b/timm/models/resnetv2.py index 43940cc3..e38eaf5e 100644 --- a/timm/models/resnetv2.py +++ b/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)