Prep a set of ResNetV2 models with GroupNorm, EvoNormB0, EvoNormS0 for BN free model experiments on TPU and IPU
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…
Reference in New Issue