def _gen_efficientnet_edge(variant,
channel_multiplier=1.0,
depth_multiplier=1.0,
pretrained=False,
**kwargs):
""" Creates an EfficientNet-EdgeTPU model
Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/edgetpu
"""
arch_def = [
# NOTE `fc` is present to override a mismatch between stem channels and in chs not
# present in other models
['er_r1_k3_s1_e4_c24_fc24_noskip'],
['er_r2_k3_s2_e8_c32'],
['er_r4_k3_s2_e8_c48'],
['ir_r5_k5_s2_e8_c96'],
['ir_r4_k5_s1_e8_c144'],
['ir_r2_k5_s2_e8_c192'],
]
round_chs_fn = partial(round_channels, multiplier=channel_multiplier)
model_kwargs = dict(
block_args=decode_arch_def(arch_def, depth_multiplier),
num_features=round_chs_fn(1280),
stem_size=32,
round_chs_fn=round_chs_fn,
norm_layer=kwargs.pop('norm_layer', None)
or partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
act_layer=resolve_act_layer(kwargs, 'prelu'),
**kwargs,
)
model = _create_effnet(variant, pretrained, **model_kwargs)
return model
def _gen_efficientnetv2_xl(variant,
channel_multiplier=1.0,
depth_multiplier=1.0,
pretrained=False,
**kwargs):
""" Creates an EfficientNet-V2 Xtra-Large model
Ref impl: https://github.com/google/automl/tree/master/efficientnetv2
Paper: `EfficientNetV2: Smaller Models and Faster Training` - https://arxiv.org/abs/2104.00298
"""
arch_def = [
['cn_r4_k3_s1_e1_c32_skip'],
['er_r8_k3_s1_e4_c64'],
['er_r8_k3_s2_e4_c96'],
['ir_r16_k3_s2_e4_c192_se0.25'],
['ir_r24_k3_s1_e6_c256_se0.25'],
['ir_r32_k3_s2_e6_c512_se0.25'],
['ir_r8_k3_s1_e6_c640_se0.25'],
]
model_kwargs = dict(
block_args=decode_arch_def(arch_def, depth_multiplier),
num_features=1280,
stem_size=32,
round_chs_fn=partial(round_channels, multiplier=channel_multiplier),
norm_layer=kwargs.pop('norm_layer', None)
or partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
act_layer=resolve_act_layer(kwargs, 'silu'),
**kwargs,
)
model = _create_effnet(variant, pretrained, **model_kwargs)
return model
def _gen_mobilenet_v2(variant,
channel_multiplier=1.0,
depth_multiplier=1.0,
fix_stem_head=False,
pretrained=False,
**kwargs):
""" Generate MobileNet-V2 network
Ref impl: https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet_v2.py
Paper: https://arxiv.org/abs/1801.04381
"""
arch_def = [
['ds_r1_k3_s1_c16'],
['ir_r2_k3_s1_e6_c24'],
['ir_r3_k3_s2_e6_c32'],
['ir_r4_k3_s2_e6_c64'],
['ir_r3_k3_s1_e6_c96'],
['ir_r3_k3_s2_e6_c160'],
['ir_r1_k3_s1_e6_c320'],
]
round_chs_fn = partial(round_channels, multiplier=channel_multiplier)
model_kwargs = dict(
block_args=decode_arch_def(arch_def,
depth_multiplier=depth_multiplier,
fix_first_last=fix_stem_head),
num_features=1280 if fix_stem_head else round_chs_fn(1280),
stem_size=32,
fix_stem=fix_stem_head,
round_chs_fn=round_chs_fn,
norm_layer=kwargs.pop('norm_layer', None)
or partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
act_layer=resolve_act_layer(kwargs, 'prelu'),
**kwargs)
model = _create_effnet(variant, pretrained, **model_kwargs)
return model
def _gen_efficientnet(variant,
channel_multiplier=1.0,
depth_multiplier=1.0,
pretrained=False,
**kwargs):
"""Creates an EfficientNet model.
Ref impl: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py
Paper: https://arxiv.org/abs/1905.11946
EfficientNet params
name: (channel_multiplier, depth_multiplier, resolution, dropout_rate)
'efficientnet-b0': (1.0, 1.0, 224, 0.2),
'efficientnet-b1': (1.0, 1.1, 240, 0.2),
'efficientnet-b2': (1.1, 1.2, 260, 0.3),
'efficientnet-b3': (1.2, 1.4, 300, 0.3),
'efficientnet-b4': (1.4, 1.8, 380, 0.4),
'efficientnet-b5': (1.6, 2.2, 456, 0.4),
'efficientnet-b6': (1.8, 2.6, 528, 0.5),
'efficientnet-b7': (2.0, 3.1, 600, 0.5),
'efficientnet-b8': (2.2, 3.6, 672, 0.5),
'efficientnet-l2': (4.3, 5.3, 800, 0.5),
Args:
channel_multiplier: multiplier to number of channels per layer
depth_multiplier: multiplier to number of repeats per stage
"""
#arch_def = [
# ['ds_r1_k3_s1_e1_c16_se0.25'],
# ['ir_r2_k3_s2_e6_c24_se0.25'],
# ['ir_r2_k5_s2_e6_c40_se0.25'],
# ['ir_r3_k3_s2_e6_c80_se0.25'],
# ['ir_r3_k5_s1_e6_c112_se0.25'],
# ['ir_r4_k5_s2_e6_c192_se0.25'],
# ['ir_r1_k3_s1_e6_c320_se0.25'],
#]
arch_def = [
['ds_r1_k3_s1_e1_c16'],
['ir_r2_k3_s1_e6_c24'],
['ir_r2_k5_s2_e6_c40'],
['ir_r3_k3_s2_e6_c80'],
['ir_r3_k5_s1_e6_c112'],
['ir_r4_k5_s2_e6_c192'],
['ir_r1_k3_s1_e6_c320'],
]
round_chs_fn = partial(round_channels, multiplier=channel_multiplier)
model_kwargs = dict(
block_args=decode_arch_def(arch_def, depth_multiplier),
num_features=round_chs_fn(1280),
stem_size=32,
round_chs_fn=round_chs_fn,
act_layer=resolve_act_layer(kwargs, 'prelu'),
norm_layer=kwargs.pop('norm_layer', None)
or partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
**kwargs,
)
model = _create_effnet(variant, pretrained, **model_kwargs)
return model
def _gen_efficientnet_lite(variant,
channel_multiplier=1.0,
depth_multiplier=1.0,
pretrained=False,
**kwargs):
"""Creates an EfficientNet-Lite model.
Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite
Paper: https://arxiv.org/abs/1905.11946
EfficientNet params
name: (channel_multiplier, depth_multiplier, resolution, dropout_rate)
'efficientnet-lite0': (1.0, 1.0, 224, 0.2),
'efficientnet-lite1': (1.0, 1.1, 240, 0.2),
'efficientnet-lite2': (1.1, 1.2, 260, 0.3),
'efficientnet-lite3': (1.2, 1.4, 280, 0.3),
'efficientnet-lite4': (1.4, 1.8, 300, 0.3),
Args:
channel_multiplier: multiplier to number of channels per layer
depth_multiplier: multiplier to number of repeats per stage
"""
arch_def = [
['ds_r1_k3_s1_e1_c16'],
['ir_r2_k3_s1_e6_c24'],
['ir_r2_k5_s2_e6_c40'],
['ir_r3_k3_s2_e6_c80'],
['ir_r3_k5_s1_e6_c112'],
['ir_r4_k5_s2_e6_c192'],
['ir_r1_k3_s1_e6_c320'],
]
model_kwargs = dict(
block_args=decode_arch_def(arch_def,
depth_multiplier,
fix_first_last=True),
num_features=1280,
stem_size=32,
fix_stem=True,
round_chs_fn=partial(round_channels, multiplier=channel_multiplier),
act_layer=resolve_act_layer(kwargs, 'prelu'),
norm_layer=kwargs.pop('norm_layer', None)
or partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
**kwargs,
)
model = _create_effnet(variant, pretrained, **model_kwargs)
return model
def _gen_efficientnetv2_s(variant,
channel_multiplier=1.0,
depth_multiplier=1.0,
rw=False,
pretrained=False,
**kwargs):
""" Creates an EfficientNet-V2 Small model
Ref impl: https://github.com/google/automl/tree/master/efficientnetv2
Paper: `EfficientNetV2: Smaller Models and Faster Training` - https://arxiv.org/abs/2104.00298
NOTE: `rw` flag sets up 'small' variant to behave like my initial v2 small model,
before ref the impl was released.
"""
arch_def = [
['cn_r2_k3_s1_e1_c24_skip'],
['er_r4_k3_s1_e4_c48'],
['er_r4_k3_s2_e4_c64'],
['ir_r6_k3_s2_e4_c128'],
['ir_r9_k3_s1_e6_c160'],
['ir_r15_k3_s2_e6_c256'],
]
num_features = 1280
if rw:
# my original variant, based on paper figure differs from the official release
arch_def[0] = ['er_r2_k3_s1_e1_c24']
arch_def[-1] = ['ir_r15_k3_s2_e6_c272']
num_features = 1792
round_chs_fn = partial(round_channels, multiplier=channel_multiplier)
model_kwargs = dict(
block_args=decode_arch_def(arch_def, depth_multiplier),
num_features=round_chs_fn(num_features),
stem_size=24,
round_chs_fn=round_chs_fn,
norm_layer=kwargs.pop('norm_layer', None)
or partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
act_layer=resolve_act_layer(kwargs, 'prelu'),
**kwargs,
)
model = _create_effnet(variant, pretrained, **model_kwargs)
return model
def _gen_efficientnet_condconv(variant,
channel_multiplier=1.0,
depth_multiplier=1.0,
experts_multiplier=1,
pretrained=False,
**kwargs):
"""Creates an EfficientNet-CondConv model.
Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/condconv
"""
arch_def = [
['ds_r1_k3_s1_e1_c16_se0.25'],
['ir_r2_k3_s2_e6_c24_se0.25'],
['ir_r2_k5_s2_e6_c40_se0.25'],
['ir_r3_k3_s2_e6_c80_se0.25'],
['ir_r3_k5_s1_e6_c112_se0.25_cc4'],
['ir_r4_k5_s2_e6_c192_se0.25_cc4'],
['ir_r1_k3_s1_e6_c320_se0.25_cc4'],
]
# NOTE unlike official impl, this one uses `cc<x>` option where x is the base number of experts for each stage and
# the expert_multiplier increases that on a per-model basis as with depth/channel multipliers
round_chs_fn = partial(round_channels, multiplier=channel_multiplier)
model_kwargs = dict(
block_args=decode_arch_def(arch_def,
depth_multiplier,
experts_multiplier=experts_multiplier),
num_features=round_chs_fn(1280),
stem_size=32,
round_chs_fn=round_chs_fn,
norm_layer=kwargs.pop('norm_layer', None)
or partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
act_layer=resolve_act_layer(kwargs, 'swish'),
**kwargs,
)
model = _create_effnet(variant, pretrained, **model_kwargs)
return model