Merge pull request #2203 from huggingface/more_mobile

Add mobilenet edgetpu defs for exp, add ol mobilenet v1 back for comp…

Author: rwightman

timm/models/_efficientnet_blocks.py

@@ -662,7 +662,7 @@ def __init__(
             mid_chs = make_divisible(force_in_chs * exp_ratio)
         else:
             mid_chs = make_divisible(in_chs * exp_ratio)
-        groups = num_groups(group_size, in_chs)
+        groups = num_groups(group_size, mid_chs)  # NOTE: Using out_chs of conv_exp for groups calc
         self.has_skip = (in_chs == out_chs and stride == 1) and not noskip
         use_aa = aa_layer is not None and stride > 1  # FIXME handle dilation
 

timm/models/_efficientnet_builder.py

@@ -221,8 +221,9 @@ def _decode_block_str(block_str):
         ))
     else:
         assert False, 'Unknown block type (%s)' % block_type
+
     if 'gs' in options:
-        block_args['group_size'] = options['gs']
+        block_args['group_size'] = int(options['gs'])
 
     return block_args, num_repeat
 

timm/models/efficientnet.py

@@ -79,6 +79,7 @@ def __init__(
             num_features: int = 1280,
             in_chans: int = 3,
             stem_size: int = 32,
+            stem_kernel_size: int = 3,
             fix_stem: bool = False,
             output_stride: int = 32,
             pad_type: str = '',
@@ -104,7 +105,7 @@ def __init__(
         # Stem
         if not fix_stem:
             stem_size = round_chs_fn(stem_size)
-        self.conv_stem = create_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type)
+        self.conv_stem = create_conv2d(in_chans, stem_size, stem_kernel_size, stride=2, padding=pad_type)
         self.bn1 = norm_act_layer(stem_size, inplace=True)
 
         # Middle stages (IR/ER/DS Blocks)
@@ -277,6 +278,7 @@ def __init__(
             feature_location: str = 'bottleneck',
             in_chans: int = 3,
             stem_size: int = 32,
+            stem_kernel_size: int = 3,
             fix_stem: bool = False,
             output_stride: int = 32,
             pad_type: str = '',
@@ -299,7 +301,7 @@ def __init__(
         # Stem
         if not fix_stem:
             stem_size = round_chs_fn(stem_size)
-        self.conv_stem = create_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type)
+        self.conv_stem = create_conv2d(in_chans, stem_size, stem_kernel_size, stride=2, padding=pad_type)
         self.bn1 = norm_act_layer(stem_size, inplace=True)
 
         # Middle stages (IR/ER/DS Blocks)
@@ -478,6 +480,34 @@ def _gen_mnasnet_small(variant, channel_multiplier=1.0, pretrained=False, **kwar
     return model
 
 
+def _gen_mobilenet_v1(
+        variant, channel_multiplier=1.0, depth_multiplier=1.0, fix_stem_head=False, pretrained=False, **kwargs):
+    """
+    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 = [
+        ['dsa_r1_k3_s1_c64'],
+        ['dsa_r2_k3_s2_c128'],
+        ['dsa_r2_k3_s2_c256'],
+        ['dsa_r6_k3_s2_c512'],
+        ['dsa_r2_k3_s2_c1024'],
+    ]
+    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=1024 if fix_stem_head else max(1024, round_chs_fn(1024)),
+        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, 'relu6'),
+        **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
@@ -1056,6 +1086,95 @@ def _gen_tinynet(
     return model
 
 
+def _gen_mobilenet_edgetpu(variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):
+    """
+    Based on definitions in: https://github.com/tensorflow/models/tree/d2427a562f401c9af118e47af2f030a0a5599f55/official/projects/edgetpu/vision
+    """
+    if 'edgetpu_v2' in variant:
+        stem_size = 64
+        stem_kernel_size = 5
+        group_size = 64
+        num_features = 1280
+        act_layer = resolve_act_layer(kwargs, 'relu')
+
+        def _arch_def(chs: List[int], group_size: int):
+            return [
+                # stage 0, 112x112 in
+                [f'cn_r1_k1_s1_c{chs[0]}'],  # NOTE with expansion==1, official impl block ends just 1x1 pwl
+                # stage 1, 112x112 in
+                [f'er_r1_k3_s2_e8_c{chs[1]}', f'er_r1_k3_s1_e4_gs{group_size}_c{chs[1]}'],
+                # stage 2, 56x56 in
+                [
+                    f'er_r1_k3_s2_e8_c{chs[2]}',
+                    f'er_r1_k3_s1_e4_gs{group_size}_c{chs[2]}',
+                    f'er_r1_k3_s1_e4_c{chs[2]}',
+                    f'er_r1_k3_s1_e4_gs{group_size}_c{chs[2]}',
+                ],
+                # stage 3, 28x28 in
+                [f'er_r1_k3_s2_e8_c{chs[3]}', f'ir_r3_k3_s1_e4_c{chs[3]}'],
+                # stage 4, 14x14in
+                [f'ir_r1_k3_s1_e8_c{chs[4]}', f'ir_r3_k3_s1_e4_c{chs[4]}'],
+                # stage 5, 14x14in
+                [f'ir_r1_k3_s2_e8_c{chs[5]}', f'ir_r3_k3_s1_e4_c{chs[5]}'],
+                # stage 6, 7x7 in
+                [f'ir_r1_k3_s1_e8_c{chs[6]}'],
+            ]
+
+        if 'edgetpu_v2_xs' in variant:
+            stem_size = 32
+            stem_kernel_size = 3
+            channels = [16, 32, 48, 96, 144, 160, 192]
+        elif 'edgetpu_v2_s' in variant:
+            channels = [24, 48, 64, 128, 160, 192, 256]
+        elif 'edgetpu_v2_m' in variant:
+            channels = [32, 64, 80, 160, 192, 240, 320]
+            num_features = 1344
+        elif 'edgetpu_v2_l' in variant:
+            stem_kernel_size = 7
+            group_size = 128
+            channels = [32, 64, 96, 192, 240, 256, 384]
+            num_features = 1408
+        else:
+            assert False
+
+        arch_def = _arch_def(channels, group_size)
+    else:
+        # v1
+        stem_size = 32
+        stem_kernel_size = 3
+        num_features = 1280
+        act_layer = resolve_act_layer(kwargs, 'relu')
+        arch_def = [
+            # stage 0, 112x112 in
+            ['cn_r1_k1_s1_c16'],
+            # stage 1, 112x112 in
+            ['er_r1_k3_s2_e8_c32', 'er_r3_k3_s1_e4_c32'],
+            # stage 2, 56x56 in
+            ['er_r1_k3_s2_e8_c48', 'er_r3_k3_s1_e4_c48'],
+            # stage 3, 28x28 in
+            ['ir_r1_k3_s2_e8_c96', 'ir_r3_k3_s1_e4_c96'],
+            # stage 4, 14x14in
+            ['ir_r1_k3_s1_e8_c96_noskip', 'ir_r3_k3_s1_e4_c96'],
+            # stage 5, 14x14in
+            ['ir_r1_k5_s2_e8_c160', 'ir_r3_k5_s1_e4_c160'],
+            # stage 6, 7x7 in
+            ['ir_r1_k3_s1_e8_c192'],
+        ]
+
+    model_kwargs = dict(
+        block_args=decode_arch_def(arch_def, depth_multiplier),
+        num_features=num_features,
+        stem_size=stem_size,
+        stem_kernel_size=stem_kernel_size,
+        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=act_layer,
+        **kwargs,
+    )
+    model = _create_effnet(variant, pretrained, **model_kwargs)
+    return model
+
+
 def _cfg(url='', **kwargs):
     return {
         'url': url, 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
@@ -1086,6 +1205,9 @@ def _cfg(url='', **kwargs):
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mnasnet_small_lamb-aff75073.pth',
         hf_hub_id='timm/'),
 
+    'mobilenet_100.untrained': _cfg(),
+    'mobilenet_125.untrained': _cfg(),
+
     'mobilenetv2_035.untrained': _cfg(),
     'mobilenetv2_050.lamb_in1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv2_050-3d30d450.pth',
@@ -1395,7 +1517,6 @@ def _cfg(url='', **kwargs):
         hf_hub_id='timm/',
         input_size=(3, 456, 456), pool_size=(15, 15), crop_pct=0.934),
 
-
     'tf_efficientnet_es.in1k': _cfg(
         url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_es-ca1afbfe.pth',
         hf_hub_id='timm/',
@@ -1584,6 +1705,23 @@ def _cfg(url='', **kwargs):
         input_size=(3, 106, 106), pool_size=(4, 4),  # int(224 * 0.475)
         url='https://github.com/huawei-noah/CV-Backbones/releases/download/v1.2.0/tinynet_e.pth',
         hf_hub_id='timm/'),
+
+    'mobilenet_edgetpu_100.untrained': _cfg(
+        # hf_hub_id='timm/',
+        input_size=(3, 224, 224), crop_pct=0.9),
+    'mobilenet_edgetpu_v2_xs.untrained': _cfg(
+        # hf_hub_id='timm/',
+        input_size=(3, 224, 224), crop_pct=0.9),
+    'mobilenet_edgetpu_v2_s.untrained': _cfg(
+        #hf_hub_id='timm/',
+        input_size=(3, 224, 224), crop_pct=0.9),
+    'mobilenet_edgetpu_v2_m.untrained': _cfg(
+        #hf_hub_id='timm/',
+        input_size=(3, 224, 224), crop_pct=0.9),
+    'mobilenet_edgetpu_v2_l.untrained': _cfg(
+        #hf_hub_id='timm/',
+        input_size=(3, 224, 224), crop_pct=0.9),
+
 })
 
 
@@ -1650,6 +1788,20 @@ def mnasnet_small(pretrained=False, **kwargs) -> EfficientNet:
     return model
 
 
+@register_model
+def mobilenet_100(pretrained=False, **kwargs) -> EfficientNet:
+    """ MobileNet V1 """
+    model = _gen_mobilenet_v1('mobilenet_100', 1.0, pretrained=pretrained, **kwargs)
+    return model
+
+
+@register_model
+def mobilenet_125(pretrained=False, **kwargs) -> EfficientNet:
+    """ MobileNet V1 """
+    model = _gen_mobilenet_v1('mobilenet_125', 1.25, pretrained=pretrained, **kwargs)
+    return model
+
+
 @register_model
 def mobilenetv2_035(pretrained=False, **kwargs) -> EfficientNet:
     """ MobileNet V2 w/ 0.35 channel multiplier """
@@ -2510,6 +2662,41 @@ def tinynet_e(pretrained=False, **kwargs) -> EfficientNet:
     return model
 
 
+@register_model
+def mobilenet_edgetpu_100(pretrained=False, **kwargs) -> EfficientNet:
+    """ MobileNet-EdgeTPU-v1 100. """
+    model = _gen_mobilenet_edgetpu('mobilenet_edgetpu_100', pretrained=pretrained, **kwargs)
+    return model
+
+
+@register_model
+def mobilenet_edgetpu_v2_xs(pretrained=False, **kwargs) -> EfficientNet:
+    """ MobileNet-EdgeTPU-v2 Extra Small. """
+    model = _gen_mobilenet_edgetpu('mobilenet_edgetpu_v2_xs', pretrained=pretrained, **kwargs)
+    return model
+
+
+@register_model
+def mobilenet_edgetpu_v2_s(pretrained=False, **kwargs) -> EfficientNet:
+    """ MobileNet-EdgeTPU-v2 Small. """
+    model = _gen_mobilenet_edgetpu('mobilenet_edgetpu_v2_s', pretrained=pretrained, **kwargs)
+    return model
+
+
+@register_model
+def mobilenet_edgetpu_v2_m(pretrained=False, **kwargs) -> EfficientNet:
+    """ MobileNet-EdgeTPU-v2 Medium. """
+    model = _gen_mobilenet_edgetpu('mobilenet_edgetpu_v2_m', pretrained=pretrained, **kwargs)
+    return model
+
+
+@register_model
+def mobilenet_edgetpu_v2_l(pretrained=False, **kwargs) -> EfficientNet:
+    """ MobileNet-EdgeTPU-v2 Large. """
+    model = _gen_mobilenet_edgetpu('mobilenet_edgetpu_v2_l', pretrained=pretrained, **kwargs)
+    return model
+
+
 register_model_deprecations(__name__, {
     'tf_efficientnet_b0_ap': 'tf_efficientnet_b0.ap_in1k',
     'tf_efficientnet_b1_ap': 'tf_efficientnet_b1.ap_in1k',