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

pull/989/head
Ross Wightman 3 years ago
parent 3819bef93e
commit af607b75cc

@ -55,7 +55,7 @@ class EvoNormBatch2d(nn.Module):
class EvoNormSample2d(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__() super(EvoNormSample2d, self).__init__()
self.apply_act = apply_act # apply activation (non-linearity) self.apply_act = apply_act # apply activation (non-linearity)
self.groups = groups self.groups = groups

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

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

Loading…
Cancel
Save