Add MobileNetV2 and MobileNetV3

README.md

@@ -7,4 +7,10 @@
 CUDA_VISIBLE_DEVICES= KERAS_BACKEND=tensorflow pytest
 ```
 
+## Work in Progress
+
+- Test pretrained weights
+- Add `MobileNet100V3SmallMinimal` pretrained weights
+- Add `MobileNet100V3LargeMinimal` pretrained weights
+
 ## Acknowledgments

kimm/blocks/base_block.py

@@ -18,8 +18,8 @@ def apply_activation(x, activation=None, name="activation"):
 
 def apply_conv2d_block(
     inputs,
-    filters,
-    kernel_size,
+    filters=None,
+    kernel_size=None,
     strides=1,
     groups=1,
     activation=None,
@@ -28,6 +28,10 @@ def apply_conv2d_block(
     bn_epsilon=1e-5,
     name="conv2d_block",
 ):
+    if kernel_size is None:
+        raise ValueError(
+            f"kernel_size must be passed. Received: kernel_size={kernel_size}"
+        )
     x = inputs
 
     padding = "same"
@@ -80,12 +84,12 @@ def apply_se_block(
     x = inputs
     x = layers.GlobalAveragePooling2D(keepdims=True, name=f"{name}_mean")(x)
     x = layers.Conv2D(
-        se_channels, 1, use_bias=True, name=f"{name}_reduce_conv2d"
+        se_channels, 1, use_bias=True, name=f"{name}_conv_reduce"
     )(x)
-    x = apply_activation(x, activation, name=f"{name}_act")
+    x = apply_activation(x, activation, name=f"{name}_act1")
     x = layers.Conv2D(
-        input_channels, 1, use_bias=True, name=f"{name}_expand_conv2d"
+        input_channels, 1, use_bias=True, name=f"{name}_conv_expand"
     )(x)
-    x = apply_activation(x, gate_activation, name=f"{name}_gate_act")
+    x = apply_activation(x, gate_activation, name=f"{name}_gate")
     out = layers.Multiply(name=name)([ori_x, x])
     return out

kimm/models/ghostnet.py

@@ -290,42 +290,40 @@ def __init__(
         x = apply_conv2d_block(
             x, stem_channels, 3, 2, activation="relu", name="conv_stem"
         )
-        features["S2"] = x
+        features["STEM_S2"] = x
 
         # blocks
         total_layer_idx = 0
-        block_idx = 0
-        net_stride = 2
-        for cfg in config:
-            layer_idx = 0
-            strides = 1
-            for kernel_size, expand_size, channels, se_ratio, strides in cfg:
-                output_channels = make_divisible(channels * width, 4)
-                hidden_channels = make_divisible(expand_size * width, 4)
+        current_stride = 2
+        for current_block_idx, cfg in enumerate(config):
+            for current_layer_idx, (k, e, c, se, s) in enumerate(cfg):
+                output_channels = make_divisible(c * width, 4)
+                hidden_channels = make_divisible(e * width, 4)
                 use_attention = False
                 if version == "v2" and total_layer_idx > 1:
                     use_attention = True
+                name = f"blocks_{current_block_idx}_{current_layer_idx}"
                 x = apply_ghost_bottleneck(
                     x,
                     hidden_channels,
                     output_channels,
-                    kernel_size,
-                    strides,
-                    se_ratio=se_ratio,
+                    k,
+                    s,
+                    se_ratio=se,
                     use_attention=use_attention,
-                    name=f"blocks{block_idx}_{layer_idx}",
+                    name=name,
                 )
-                layer_idx += 1
                 total_layer_idx += 1
-            if strides > 1:
-                net_stride *= strides
-                # add feature
-                features[f"S{net_stride}"] = x
-            block_idx += 1
+            current_stride *= s
+            features[f"BLOCK{current_block_idx}_S{current_stride}"] = x
         # post stages conv block
-        output_channels = make_divisible(expand_size * width, 4)
+        output_channels = make_divisible(e * width, 4)
         x = apply_conv2d_block(
-            x, output_channels, 1, activation="relu", name=f"blocks{block_idx}"
+            x,
+            output_channels,
+            1,
+            activation="relu",
+            name=f"blocks_{current_block_idx+1}",
         )
 
         if include_top:
@@ -366,8 +364,14 @@ def __init__(
 
     @staticmethod
     def available_feature_keys():
-        # predefined for better UX
-        return [f"S{2**i}" for i in range(1, 6)]
+        feature_keys = ["STEM_S2"]
+        feature_keys.extend(
+            [
+                f"BLOCK{i}_S{j}"
+                for i, j in zip(range(9), [2, 4, 4, 8, 8, 16, 16, 32, 32])
+            ]
+        )
+        return feature_keys
 
     def get_config(self):
         config = super().get_config()

kimm/models/ghostnet_test.py

@@ -25,11 +25,14 @@ def test_ghostnet_feature_extractor(self, model_class):
         y = model.predict(x)
 
         self.assertIsInstance(y, dict)
-        self.assertEqual(list(y["S2"].shape), [1, 112, 112, 16])
-        self.assertEqual(list(y["S4"].shape), [1, 56, 56, 24])
-        self.assertEqual(list(y["S8"].shape), [1, 28, 28, 40])
-        self.assertEqual(list(y["S16"].shape), [1, 14, 14, 80])
-        self.assertEqual(list(y["S32"].shape), [1, 7, 7, 160])
+        self.assertAllEqual(
+            list(y.keys()), model_class.available_feature_keys()
+        )
+        self.assertEqual(list(y["STEM_S2"].shape), [1, 112, 112, 16])
+        self.assertEqual(list(y["BLOCK1_S4"].shape), [1, 56, 56, 24])
+        self.assertEqual(list(y["BLOCK3_S8"].shape), [1, 28, 28, 40])
+        self.assertEqual(list(y["BLOCK5_S16"].shape), [1, 14, 14, 80])
+        self.assertEqual(list(y["BLOCK7_S32"].shape), [1, 7, 7, 160])
 
     @parameterized.named_parameters([(GhostNet100V2.__name__, GhostNet100V2)])
     def test_ghostnetv2_base(self, model_class):
@@ -49,8 +52,11 @@ def test_ghostnetv2_feature_extractor(self, model_class):
         y = model.predict(x)
 
         self.assertIsInstance(y, dict)
-        self.assertEqual(list(y["S2"].shape), [1, 112, 112, 16])
-        self.assertEqual(list(y["S4"].shape), [1, 56, 56, 24])
-        self.assertEqual(list(y["S8"].shape), [1, 28, 28, 40])
-        self.assertEqual(list(y["S16"].shape), [1, 14, 14, 80])
-        self.assertEqual(list(y["S32"].shape), [1, 7, 7, 160])
+        self.assertAllEqual(
+            list(y.keys()), model_class.available_feature_keys()
+        )
+        self.assertEqual(list(y["STEM_S2"].shape), [1, 112, 112, 16])
+        self.assertEqual(list(y["BLOCK1_S4"].shape), [1, 56, 56, 24])
+        self.assertEqual(list(y["BLOCK3_S8"].shape), [1, 28, 28, 40])
+        self.assertEqual(list(y["BLOCK5_S16"].shape), [1, 14, 14, 80])
+        self.assertEqual(list(y["BLOCK7_S32"].shape), [1, 7, 7, 160])