Add efficientnet family including efficient, efficientlite and effici…

…entv2 (#4)

* Fix tf pretrained model

* Code cleanup

* Add comments

* Add draft efficientnet

* Add unit test for efficientnet

* Add efficientnet lite

* Improve unit tests

* Merge efficientnet lite into efficientnet

* Add efficientnet v2

README.md

@@ -10,7 +10,5 @@ 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

@@ -24,20 +24,25 @@ def apply_conv2d_block(
     groups=1,
     activation=None,
     use_depthwise=False,
+    add_skip=False,
     bn_momentum=0.9,
     bn_epsilon=1e-5,
+    padding=None,
     name="conv2d_block",
 ):
     if kernel_size is None:
         raise ValueError(
             f"kernel_size must be passed. Received: kernel_size={kernel_size}"
         )
+    input_channels = inputs.shape[-1]
+    has_skip = add_skip and strides == 1 and input_channels == filters
     x = inputs
 
-    padding = "same"
-    if strides > 1:
-        padding = "valid"
-        x = layers.ZeroPadding2D(kernel_size // 2, name=f"{name}_pad")(x)
+    if padding is None:
+        padding = "same"
+        if strides > 1:
+            padding = "valid"
+            x = layers.ZeroPadding2D(kernel_size // 2, name=f"{name}_pad")(x)
 
     if not use_depthwise:
         x = layers.Conv2D(
@@ -61,6 +66,8 @@ def apply_conv2d_block(
         name=f"{name}_bn", momentum=bn_momentum, epsilon=bn_epsilon
     )(x)
     x = apply_activation(x, activation, name=name)
+    if has_skip:
+        x = layers.Add()([x, inputs])
     return x
 
 
@@ -70,14 +77,17 @@ def apply_se_block(
     activation="relu",
     gate_activation="sigmoid",
     make_divisible_number=None,
+    se_input_channels=None,
     name="se_block",
 ):
     input_channels = inputs.shape[-1]
+    if se_input_channels is None:
+        se_input_channels = input_channels
     if make_divisible_number is None:
-        se_channels = round(input_channels * se_ratio)
+        se_channels = round(se_input_channels * se_ratio)
     else:
         se_channels = make_divisible(
-            input_channels * se_ratio, make_divisible_number
+            se_input_channels * se_ratio, make_divisible_number
         )
 
     ori_x = inputs

kimm/layers/attention.py

@@ -26,23 +26,53 @@ def __init__(
         self.name = name
 
         self.qkv = layers.Dense(
-            hidden_dim * 3, use_bias=use_qkv_bias, name=f"{name}_qkv"
+            hidden_dim * 3,
+            use_bias=use_qkv_bias,
+            dtype=self.dtype,
+            name=f"{name}_qkv",
         )
         if use_qk_norm:
-            self.q_norm = layers.LayerNormalization(name=f"{name}_q_norm")
-            self.k_norm = layers.LayerNormalization(name=f"{name}_k_norm")
+            self.q_norm = layers.LayerNormalization(
+                dtype=self.dtype_policy, name=f"{name}_q_norm"
+            )
+            self.k_norm = layers.LayerNormalization(
+                dtype=self.dtype_policy, name=f"{name}_k_norm"
+            )
         else:
-            self.q_norm = layers.Identity()
-            self.k_norm = layers.Identity()
+            self.q_norm = layers.Identity(dtype=self.dtype_policy)
+            self.k_norm = layers.Identity(dtype=self.dtype_policy)
 
         self.attention_dropout = layers.Dropout(
-            attention_dropout_rate, name=f"{name}_attn_drop"
+            attention_dropout_rate,
+            dtype=self.dtype_policy,
+            name=f"{name}_attn_drop",
+        )
+        self.projection = layers.Dense(
+            hidden_dim, dtype=self.dtype_policy, name=f"{name}_proj"
         )
-        self.projection = layers.Dense(hidden_dim, name=f"{name}_proj")
         self.projection_dropout = layers.Dropout(
-            projection_dropout_rate, name=f"{name}_proj_drop"
+            projection_dropout_rate,
+            dtype=self.dtype_policy,
+            name=f"{name}_proj_drop",
         )
 
+    def build(self, input_shape):
+        self.qkv.build(input_shape)
+        qkv_output_shape = list(input_shape)
+        qkv_output_shape[-1] = qkv_output_shape[-1] * 3
+        self.q_norm.build(qkv_output_shape)
+        self.k_norm.build(qkv_output_shape)
+        attention_input_shape = [
+            input_shape[0],
+            self.num_heads,
+            input_shape[1],
+            input_shape[1],
+        ]
+        self.attention_dropout.build(attention_input_shape)
+        self.projection.build(input_shape)
+        self.projection_dropout.build(input_shape)
+        self.built = True
+
     def call(self, inputs, training=None, mask=None):
         input_shape = ops.shape(inputs)
         qkv = self.qkv(inputs)

kimm/layers/attention_test.py

@@ -1,15 +1,18 @@
 from absl.testing import parameterized
-from keras import random
 from keras.src import testing
 
 from kimm.layers.attention import Attention
 
 
 class AttentionTest(testing.TestCase, parameterized.TestCase):
-    def test_attention(self):
-        x = random.uniform([1, 197, 768])
-        layer = Attention(768)
-
-        y = layer(x)
-
-        self.assertEqual(y.shape, [1, 197, 768])
+    def test_attention_basic(self):
+        self.run_layer_test(
+            Attention,
+            init_kwargs={"hidden_dim": 20, "num_heads": 2},
+            input_shape=(1, 10, 20),
+            expected_output_shape=(1, 10, 20),
+            expected_num_trainable_weights=3,
+            expected_num_non_trainable_weights=0,
+            expected_num_losses=0,
+            supports_masking=False,
+        )

kimm/layers/layer_scale.py

@@ -30,7 +30,7 @@ def get_config(self):
         config.update(
             {
                 "hidden_size": self.hidden_size,
-                "initializer": self.initializer,
+                "initializer": initializers.serialize(self.initializer),
                 "name": self.name,
             }
         )

kimm/layers/layer_scale_test.py

@@ -1,15 +1,18 @@
 from absl.testing import parameterized
-from keras import random
 from keras.src import testing
 
 from kimm.layers.layer_scale import LayerScale
 
 
 class LayerScaleTest(testing.TestCase, parameterized.TestCase):
-    def test_layer_scale(self):
-        x = random.uniform([1, 123])
-        layer = LayerScale(123)
-
-        y = layer(x)
-
-        self.assertEqual(y.shape, [1, 123])
+    def test_layer_scale_basic(self):
+        self.run_layer_test(
+            LayerScale,
+            init_kwargs={"hidden_size": 10},
+            input_shape=(1, 10),
+            expected_output_shape=(1, 10),
+            expected_num_trainable_weights=1,
+            expected_num_non_trainable_weights=0,
+            expected_num_losses=0,
+            supports_masking=False,
+        )

kimm/layers/position_embedding.py

@@ -7,6 +7,11 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     def build(self, input_shape):
+        if len(input_shape) != 3:
+            raise ValueError(
+                "PositionEmbedding only accepts 3-dimensional input. "
+                f"Received: input_shape={input_shape}"
+            )
         self.pos_embed = self.add_weight(
             shape=[1, input_shape[-2] + 1, input_shape[-1]],
             initializer="random_normal",
@@ -26,6 +31,11 @@ def call(self, inputs, training=None, mask=None):
         x = ops.add(x, self.pos_embed)
         return x
 
+    def compute_output_shape(self, input_shape):
+        output_shape = list(input_shape)
+        output_shape[1] = output_shape[1] + 1
+        return output_shape
+
     def get_config(self):
         return super().get_config()
 

kimm/layers/position_embedding_test.py

@@ -1,15 +1,26 @@
 from absl.testing import parameterized
-from keras import random
+from keras import layers
 from keras.src import testing
 
 from kimm.layers.position_embedding import PositionEmbedding
 
 
 class PositionEmbeddingTest(testing.TestCase, parameterized.TestCase):
-    def test_position_embedding(self):
-        x = random.uniform([1, 123, 768])
-        layer = PositionEmbedding()
+    def test_position_embedding_basic(self):
+        self.run_layer_test(
+            PositionEmbedding,
+            init_kwargs={},
+            input_shape=(1, 10, 10),
+            expected_output_shape=(1, 11, 10),
+            expected_num_trainable_weights=2,
+            expected_num_non_trainable_weights=0,
+            expected_num_losses=0,
+            supports_masking=False,
+        )
 
-        y = layer(x)
-
-        self.assertEqual(y.shape, [1, 124, 768])
+    def test_position_embedding_invalid_input_shape(self):
+        inputs = layers.Input([3])
+        with self.assertRaisesRegex(
+            ValueError, "PositionEmbedding only accepts 3-dimensional input."
+        ):
+            PositionEmbedding()(inputs)