Repository: james77777778/keras-image-models (language: python)

Issue:
Add popular models
Reference:
- https://huggingface.co/timm?sort_models=downloads#models
- https://huggingface.co/collections/timm/timm-takes-on-the-classics-655ce63dcaa21067cd7d3f4c

Classics:
- [x] densenet
- [x] inceptionv3
- [ ] inceptionresnetv2
- [x] resnet
- [ ] squeezenet
- [x] vgg
- [x] xception

Lightweight
- [ ] edgenext
- [x] efficientnetlite
- [ ] efficientvit
- [x] ghostnet
- [x] ghostnetv2
- [x] ghostnetv3
- [x] lcnet
- [ ] levit
- [x] mobilenetv2
- [x] mobilenetv3
- [x] mobileone
- [x] mobilevit
- [x] mobilevitv2
- [ ] repghostnet
- [x] tinynet

SOTA:
- [x] convmixer
- [x] convnext
- [x] efficientnet
- [x] efficientnetv2
- [ ] eva02
- [x] hgnet
- [x] hgnetv2
- [x] inceptionnext
- [ ] maxvit
- [x] regnet
- [x] repvgg
- [ ] swin transformer
- [x] vit


Discussion / hints:

Fix patch:
diff --git a/kimm/models/__init__.py b/kimm/models/__init__.py
index d1501e5..bbd825a 100644
--- a/kimm/models/__init__.py
+++ b/kimm/models/__init__.py
@@ -1,4 +1,5 @@
 from kimm.models.base_model import BaseModel
+from kimm.models.convmixer import *  # noqa:F403
 from kimm.models.densenet import *  # noqa:F403
 from kimm.models.efficientnet import *  # noqa:F403
 from kimm.models.ghostnet import *  # noqa:F403
diff --git a/kimm/models/convmixer.py b/kimm/models/convmixer.py
new file mode 100644
index 0000000..76ee965
--- /dev/null
+++ b/kimm/models/convmixer.py
@@ -0,0 +1,285 @@
+import typing
+
+import keras
+from keras import layers
+from keras import utils
+
+from kimm.models.base_model import BaseModel
+from kimm.utils import add_model_to_registry
+
+
+def apply_convmixer_block(
+    inputs, output_channels, kernel_size, activation, name="convmixer_block"
+):
+    x = inputs
+
+    # Depthwise
+    x = layers.DepthwiseConv2D(
+        kernel_size,
+        1,
+        padding="same",
+        activation=activation,
+        use_bias=True,
+        name=f"{name}_0_fn_0_dwconv2d",
+    )(x)
+    x = layers.BatchNormalization(
+        momentum=0.9, epsilon=1e-5, name=f"{name}_0_fn_2"
+    )(x)
+    x = layers.Add()([x, inputs])
+
+    # Pointwise
+    x = layers.Conv2D(
+        output_channels,
+        1,
+        1,
+        activation=activation,
+        use_bias=True,
+        name=f"{name}_1_conv2d",
+    )(x)
+    x = layers.BatchNormalization(momentum=0.9, epsilon=1e-5, name=f"{name}_3")(
+        x
+    )
+    return x
+
+
+class ConvMixer(BaseModel):
+    def __init__(
+        self,
+        depth: int = 32,
+        hidden_channels: int = 768,
+        patch_size: int = 7,
+        kernel_size: int = 7,
+        activation: str = "relu",
+        **kwargs,
+    ):
+        parsed_kwargs = self.parse_kwargs(kwargs)
+        img_input = self.determine_input_tensor(
+            parsed_kwargs["input_tensor"],
+            parsed_kwargs["input_shape"],
+            parsed_kwargs["default_size"],
+        )
+        x = img_input
+
+        if parsed_kwargs["include_preprocessing"]:
+            x = self.build_preprocessing(x, "imagenet")
+
+        # Prepare feature extraction
+        features = {}
+
+        # Stem
+        x = layers.Conv2D(
+            hidden_channels,
+            patch_size,
+            patch_size,
+            activation=activation,
+            use_bias=True,
+            name="stem_conv2d",
+        )(x)
+        x = layers.BatchNormalization(
+            momentum=0.9, epsilon=1e-5, name="stem_bn"
+        )(x)
+        features["STEM"] = x
+
+        # Blocks
+        for i in range(depth):
+            x = apply_convmixer_block(
+                x, hidden_channels, kernel_size, activation, name=f"blocks_{i}"
+            )
+            # Add feature
+            features[f"BLOCK{i}"] = x
+
+        # Head
+        if parsed_kwargs["include_top"]:
+            x = self.build_top(
+                x,
+                parsed_kwargs["classes"],
+                parsed_kwargs["classifier_activation"],
+                parsed_kwargs["dropout_rate"],
+            )
+        else:
+            if parsed_kwargs["pooling"] == "avg":
+                x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
+            elif parsed_kwargs["pooling"] == "max":
+                x = layers.GlobalMaxPooling2D(name="max_pool")(x)
+
+        # Ensure that the model takes into account
+        # any potential predecessors of `input_tensor`.
+        if parsed_kwargs["input_tensor"] is not None:
+            inputs = utils.get_source_inputs(parsed_kwargs["input_tensor"])
+        else:
+            inputs = img_input
+
+        super().__init__(inputs=inputs, outputs=x, features=features, **kwargs)
+
+        # All references to `self` below this line
+        self.add_references(parsed_kwargs)
+        self.depth = depth
+        self.hidden_channels = hidden_channels
+        self.patch_size = patch_size
+        self.kernel_size = kernel_size
+        self.activation = activation
+
+    @staticmethod
+    def available_feature_keys():
+        raise NotImplementedError
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "depth": self.depth,
+                "hidden_channels": self.hidden_channels,
+                "patch_size": self.patch_size,
+                "kernel_size": self.kernel_size,
+                "activation": self.activation,
+            }
+        )
+        return config
+
+    def fix_config(self, config):
+        unused_kwargs = [
+            "depth",
+            "hidden_channels",
+            "patch_size",
+            "kernel_size",
+            "activation",
+        ]
+        for k in unused_kwargs:
+            config.pop(k, None)
+        return config
+
+
+"""
+Model Definition
+"""
+
+
+class ConvMixer736D32(ConvMixer):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,
+        name: str = "ConvMixer736D32",
+        **kwargs,
+    ):
+        kwargs = self.fix_config(kwargs)
+        super().__init__(
+            32,
+            768,
+            7,
+            7,
+            "relu",
+            input_tensor=input_tensor,
+            input_shape=input_shape,
+            include_preprocessing=include_preprocessing,
+            include_top=include_top,
+            pooling=pooling,
+            dropout_rate=dropout_rate,
+            classes=classes,
+            classifier_activation=classifier_activation,
+            weights=weights,
+            name=name,
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM"]
+        feature_keys.extend([f"BLOCK{i}" for i in range(32)])
+        return feature_keys
+
+
+class ConvMixer1024D20(ConvMixer):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,
+        name: str = "ConvMixer1024D20",
+        **kwargs,
+    ):
+        kwargs = self.fix_config(kwargs)
+        super().__init__(
+            20,
+            1024,
+            14,
+            9,
+            "gelu",
+            input_tensor=input_tensor,
+            input_shape=input_shape,
+            include_preprocessing=include_preprocessing,
+            include_top=include_top,
+            pooling=pooling,
+            dropout_rate=dropout_rate,
+            classes=classes,
+            classifier_activation=classifier_activation,
+            weights=weights,
+            name=name,
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM"]
+        feature_keys.extend([f"BLOCK{i}" for i in range(20)])
+        return feature_keys
+
+
+class ConvMixer1536D20(ConvMixer):
+    def __init__(
+        self,
+        input_tensor: keras.KerasTensor = None,
+        input_shape: typing.Optional[typing.Sequence[int]] = None,
+        include_preprocessing: bool = True,
+        include_top: bool = True,
+        pooling: typing.Optional[str] = None,
+        dropout_rate: float = 0.0,
+        classes: int = 1000,
+        classifier_activation: str = "softmax",
+        weights: typing.Optional[str] = None,
+        name: str = "ConvMixer1536D20",
+        **kwargs,
+    ):
+        kwargs = self.fix_config(kwargs)
+        super().__init__(
+            20,
+            1536,
+            7,
+            9,
+            "gelu",
+            input_tensor=input_tensor,
+            input_shape=input_shape,
+            include_preprocessing=include_preprocessing,
+            include_top=include_top,
+            pooling=pooling,
+            dropout_rate=dropout_rate,
+            classes=classes,
+            classifier_activation=classifier_activation,
+            weights=weights,
+            name=name,
+            **kwargs,
+        )
+
+    @staticmethod
+    def available_feature_keys():
+        feature_keys = ["STEM"]
+        feature_keys.extend([f"BLOCK{i}" for i in range(20)])
+        return feature_keys
+
+
+add_model_to_registry(ConvMixer736D32, "imagenet")
+add_model_to_registry(ConvMixer1024D20, "imagenet")
+add_model_to_registry(ConvMixer1536D20, "imagenet")
diff --git a/shell/export.sh b/shell/export.sh
index bdafae8..ef849b2 100755
--- a/shell/export.sh
+++ b/shell/export.sh
@@ -4,6 +4,7 @@ set -Euxo pipefail
 export CUDA_VISIBLE_DEVICES=
 export TF_CPP_MIN_LOG_LEVEL=3
 export KERAS_BACKEND=tensorflow
+python3 -m tools.convert_convmixer_from_timm
 python3 -m tools.convert_densenet_from_timm
 python3 -m tools.convert_efficientnet_from_timm
 python3 -m tools.convert_ghostnet_from_timm
diff --git a/tools/convert_convmixer_from_timm.py b/tools/convert_convmixer_from_timm.py
new file mode 100644
index 0000000..cbff78a
--- /dev/null
+++ b/tools/convert_convmixer_from_timm.py
@@ -0,0 +1,123 @@
+"""
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+pip install timm
+"""
+import os
+
+import keras
+import numpy as np
+import timm
+import torch
+
+from kimm.models import convmixer
+from kimm.utils.timm_utils import assign_weights
+from kimm.utils.timm_utils import is_same_weights
+from kimm.utils.timm_utils import separate_keras_weights
+from kimm.utils.timm_utils import separate_torch_state_dict
+
+timm_model_names = [
+    "convmixer_768_32.in1k",
+    "convmixer_1024_20_ks9_p14.in1k",
+    "convmixer_1536_20.in1k",
+]
+keras_model_classes = [
+    convmixer.ConvMixer736D32,
+    convmixer.ConvMixer1024D20,
+    convmixer.ConvMixer1536D20,
+]
+
+for timm_model_name, keras_model_class in zip(
+    timm_model_names, keras_model_classes
+):
+    """
+    Prepare timm model and keras model
+    """
+    input_shape = [224, 224, 3]
+    torch_model = timm.create_model(timm_model_name, pretrained=True)
+    torch_model = torch_model.eval()
+    trainable_state_dict, non_trainable_state_dict = separate_torch_state_dict(
+        torch_model.state_dict()
+    )
+    keras_model = keras_model_class(
+        input_shape=input_shape,
+        include_preprocessing=False,
+        classifier_activation="linear",
+    )
+    trainable_weights, non_trainable_weights = separate_keras_weights(
+        keras_model
+    )
+
+    # for torch_name, (_, keras_name) in zip(
+    #     trainable_state_dict.keys(), trainable_weights
+    # ):
+    #     print(f"{torch_name}    {keras_name}")
+
+    # print(len(trainable_state_dict.keys()))
+    # print(len(trainable_weights))
+
+    # exit()
+
+    """
+    Assign weights
+    """
+    for keras_weight, keras_name in trainable_weights + non_trainable_weights:
+        keras_name: str
+        torch_name = keras_name
+        torch_name = torch_name.replace("_", ".")
+        # stem
+        torch_name = torch_name.replace("stem.conv2d", "stem.0")
+        torch_name = torch_name.replace("stem.bn", "stem.2")
+        # blocks
+        torch_name = torch_name.replace("dwconv2d.", "")
+        torch_name = torch_name.replace("conv2d.", "")
+        # head
+        torch_name = torch_name.replace("classifier", "head")
+
+        # weights naming mapping
+        torch_name = torch_name.replace("kernel", "weight")  # conv2d
+        torch_name = torch_name.replace("gamma", "weight")  # bn
+        torch_name = torch_name.replace("beta", "bias")  # bn
+        torch_name = torch_name.replace("moving.mean", "running_mean")  # bn
+        torch_name = torch_name.replace("moving.variance", "running_var")  # bn
+
+        # assign weights
+        if torch_name in trainable_state_dict:
+            torch_weights = trainable_state_dict[torch_name].numpy()
+        elif torch_name in non_trainable_state_dict:
+            torch_weights = non_trainable_state_dict[torch_name].numpy()
+        else:
+            raise ValueError(
+                "Can't find the corresponding torch weights. "
+                f"Got keras_name={keras_name}, torch_name={torch_name}"
+            )
+        if is_same_weights(keras_name, keras_weight, torch_name, torch_weights):
+            assign_weights(keras_name, keras_weight, torch_weights)
+        else:
+            raise ValueError(
+                "Can't find the corresponding torch weights. The shape is "
+                f"mismatched. Got keras_name={keras_name}, "
+                f"keras_weight shape={keras_weight.shape}, "
+                f"torch_name={torch_name}, "
+                f"torch_weights shape={torch_weights.shape}"
+            )
+
+    """
+    Verify model outputs
+    """
+    np.random.seed(2023)
+    keras_data = np.random.uniform(size=[1] + input_shape).astype("float32")
+    torch_data = torch.from_numpy(np.transpose(keras_data, [0, 3, 1, 2]))
+    torch_y = torch_model(torch_data)
+    keras_y = keras_model(keras_data, training=False)
+    torch_y = torch_y.detach().cpu().numpy()
+    keras_y = keras.ops.convert_to_numpy(keras_y)
+    np.testing.assert_allclose(torch_y, keras_y, atol=1e-4)
+    print(f"{keras_model_class.__name__}: output matched!")
+
+    """
+    Save converted model
+    """
+    os.makedirs("exported", exist_ok=True)
+    export_path = f"exported/{keras_model.name.lower()}_{timm_model_name}.keras"
+    keras_model.save(export_path)
+    print(f"Export to {export_path}")
diff --git a/tools/convert_xception_from_keras.py b/tools/convert_xception_from_keras.py
index 7f49ca9..77dd946 100644
--- a/tools/convert_xception_from_keras.py
+++ b/tools/convert_xception_from_keras.py
@@ -1,7 +1,3 @@
-"""
-pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
-pip install timm
-"""
 import os
 import tempfile
 


Test patch:
diff --git a/kimm/models/densenet_test.py b/kimm/models/densenet_test.py
deleted file mode 100644
index db07ddf..0000000
--- a/kimm/models/densenet_test.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.densenet import DenseNet121
-
-
-class DenseNetTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters([(DenseNet121.__name__, DenseNet121)])
-    def test_densenet_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(input_shape=[224, 224, 3])
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters([(DenseNet121.__name__, DenseNet121)])
-    def test_densenet_feature_extractor(self, model_class):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(input_shape=[224, 224, 3], feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(list(y["STEM_S4"].shape), [1, 56, 56, 64])
-        self.assertEqual(list(y["BLOCK0_S8"].shape), [1, 28, 28, 128])
-        self.assertEqual(list(y["BLOCK1_S16"].shape), [1, 14, 14, 256])
-        self.assertEqual(list(y["BLOCK2_S32"].shape), [1, 7, 7, 512])
-        self.assertEqual(list(y["BLOCK3_S32"].shape), [1, 7, 7, 1024])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters([(DenseNet121.__name__, DenseNet121, 224)])
-    def test_densenet_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class(input_shape=[224, 224, 3])
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/efficientnet_test.py b/kimm/models/efficientnet_test.py
deleted file mode 100644
index 55eb062..0000000
--- a/kimm/models/efficientnet_test.py
+++ /dev/null
@@ -1,144 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.efficientnet import EfficientNetB0
-from kimm.models.efficientnet import EfficientNetB2
-from kimm.models.efficientnet import EfficientNetLiteB0
-from kimm.models.efficientnet import EfficientNetLiteB2
-from kimm.models.efficientnet import EfficientNetV2B0
-from kimm.models.efficientnet import EfficientNetV2S
-from kimm.models.efficientnet import TinyNetE
-from kimm.utils import make_divisible
-
-
-class EfficientNetTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        [
-            (EfficientNetB0.__name__, EfficientNetB0, 224),
-            (EfficientNetB2.__name__, EfficientNetB2, 260),
-            (EfficientNetLiteB0.__name__, EfficientNetLiteB0, 224),
-            (EfficientNetLiteB2.__name__, EfficientNetLiteB2, 260),
-            (EfficientNetV2S.__name__, EfficientNetV2S, 300),
-            (EfficientNetV2B0.__name__, EfficientNetV2B0, 192),
-            (TinyNetE.__name__, TinyNetE, 106),
-        ]
-    )
-    def test_efficentnet_base(self, model_class, image_size):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters(
-        [
-            (EfficientNetB0.__name__, EfficientNetB0, 1.0, False),
-            (EfficientNetB2.__name__, EfficientNetB2, 1.1, False),
-            (EfficientNetLiteB0.__name__, EfficientNetLiteB0, 1.0, False),
-            (EfficientNetLiteB2.__name__, EfficientNetLiteB2, 1.1, False),
-            (TinyNetE.__name__, TinyNetE, 0.51, True),
-        ]
-    )
-    def test_efficentnet_feature_extractor(
-        self, model_class, width, fix_stem_channels
-    ):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(input_shape=[224, 224, 3], feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        if fix_stem_channels:
-            self.assertEqual(list(y["STEM_S2"].shape), [1, 112, 112, 32])
-        else:
-            self.assertEqual(
-                list(y["STEM_S2"].shape),
-                [1, 112, 112, make_divisible(32 * width)],
-            )
-        self.assertEqual(
-            list(y["BLOCK1_S4"].shape), [1, 56, 56, make_divisible(24 * width)]
-        )
-        self.assertEqual(
-            list(y["BLOCK2_S8"].shape), [1, 28, 28, make_divisible(40 * width)]
-        )
-        self.assertEqual(
-            list(y["BLOCK3_S16"].shape), [1, 14, 14, make_divisible(80 * width)]
-        )
-        self.assertEqual(
-            list(y["BLOCK5_S32"].shape), [1, 7, 7, make_divisible(192 * width)]
-        )
-
-    @parameterized.named_parameters(
-        [
-            (EfficientNetV2S.__name__, EfficientNetV2S, 1.0),
-            (EfficientNetV2B0.__name__, EfficientNetV2B0, 1.0),
-        ]
-    )
-    def test_efficentnet_v2_feature_extractor(self, model_class, width):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(input_shape=[224, 224, 3], feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        if "EfficientNetV2S" in model_class.__name__:
-            self.assertEqual(list(y["STEM_S2"].shape), [1, 112, 112, 24])
-            self.assertEqual(list(y["BLOCK1_S4"].shape), [1, 56, 56, 48])
-            self.assertEqual(list(y["BLOCK2_S8"].shape), [1, 28, 28, 64])
-            self.assertEqual(list(y["BLOCK3_S16"].shape), [1, 14, 14, 128])
-            self.assertEqual(list(y["BLOCK5_S32"].shape), [1, 7, 7, 256])
-        elif "EfficientNetV2B" in model_class.__name__:
-            self.assertEqual(
-                list(y["STEM_S2"].shape),
-                [1, 112, 112, make_divisible(32 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK1_S4"].shape),
-                [1, 56, 56, make_divisible(32 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK2_S8"].shape),
-                [1, 28, 28, make_divisible(48 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK3_S16"].shape),
-                [1, 14, 14, make_divisible(96 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK5_S32"].shape),
-                [1, 7, 7, make_divisible(192 * width)],
-            )
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [
-            (EfficientNetB0.__name__, EfficientNetB0, 224),
-            (EfficientNetLiteB0.__name__, EfficientNetLiteB0, 224),
-            (TinyNetE.__name__, TinyNetE, 106),
-            (EfficientNetV2S.__name__, EfficientNetV2S, 300),
-        ]
-    )
-    def test_efficientnet_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/ghostnet_test.py b/kimm/models/ghostnet_test.py
deleted file mode 100644
index c70c4fa..0000000
--- a/kimm/models/ghostnet_test.py
+++ /dev/null
@@ -1,85 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.ghostnet import GhostNet100
-from kimm.models.ghostnet import GhostNet100V2
-
-
-class GhostNetTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters([(GhostNet100.__name__, GhostNet100)])
-    def test_ghostnet_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters([(GhostNet100.__name__, GhostNet100)])
-    def test_ghostnet_feature_extractor(self, model_class):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.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):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters([(GhostNet100V2.__name__, GhostNet100V2)])
-    def test_ghostnetv2_feature_extractor(self, model_class):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.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])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [
-            (GhostNet100.__name__, GhostNet100, 224),
-            (GhostNet100V2.__name__, GhostNet100V2, 224),
-        ]
-    )
-    def test_ghostnet_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/inception_v3_test.py b/kimm/models/inception_v3_test.py
deleted file mode 100644
index e098557..0000000
--- a/kimm/models/inception_v3_test.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.inception_v3 import InceptionV3
-
-
-class InceptionV3Test(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters([(InceptionV3.__name__, InceptionV3)])
-    def test_inception_v3_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 299, 299, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters([(InceptionV3.__name__, InceptionV3)])
-    def test_inception_v3_feature_extractor(self, model_class):
-        x = random.uniform([1, 299, 299, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(list(y["STEM_S2"].shape), [1, 147, 147, 64])
-        self.assertEqual(list(y["BLOCK0_S4"].shape), [1, 71, 71, 192])
-        self.assertEqual(list(y["BLOCK1_S8"].shape), [1, 35, 35, 288])
-        self.assertEqual(list(y["BLOCK2_S16"].shape), [1, 17, 17, 768])
-        self.assertEqual(list(y["BLOCK3_S32"].shape), [1, 8, 8, 2048])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters([(InceptionV3.__name__, InceptionV3, 299)])
-    def test_inception_v3_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/mobilenet_v2_test.py b/kimm/models/mobilenet_v2_test.py
deleted file mode 100644
index 27310f9..0000000
--- a/kimm/models/mobilenet_v2_test.py
+++ /dev/null
@@ -1,75 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.mobilenet_v2 import MobileNet050V2
-from kimm.models.mobilenet_v2 import MobileNet100V2
-from kimm.utils import make_divisible
-
-
-class MobileNetV2Test(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        [
-            (MobileNet050V2.__name__, MobileNet050V2),
-            (MobileNet100V2.__name__, MobileNet100V2),
-        ]
-    )
-    def test_mobilenet_v2_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters(
-        [
-            (MobileNet050V2.__name__, MobileNet050V2, 0.5),
-            (MobileNet100V2.__name__, MobileNet100V2, 1.0),
-        ]
-    )
-    def test_mobilenet_v2_feature_extractor(self, model_class, width):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(
-            list(y["STEM_S2"].shape), [1, 112, 112, make_divisible(32 * width)]
-        )
-        self.assertEqual(
-            list(y["BLOCK1_S4"].shape), [1, 56, 56, make_divisible(24 * width)]
-        )
-        self.assertEqual(
-            list(y["BLOCK2_S8"].shape), [1, 28, 28, make_divisible(32 * width)]
-        )
-        self.assertEqual(
-            list(y["BLOCK3_S16"].shape), [1, 14, 14, make_divisible(64 * width)]
-        )
-        self.assertEqual(
-            list(y["BLOCK5_S32"].shape), [1, 7, 7, make_divisible(160 * width)]
-        )
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [(MobileNet050V2.__name__, MobileNet050V2, 224)]
-    )
-    def test_mobilenet_v2_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/mobilenet_v3_test.py b/kimm/models/mobilenet_v3_test.py
deleted file mode 100644
index c3c6e8d..0000000
--- a/kimm/models/mobilenet_v3_test.py
+++ /dev/null
@@ -1,153 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.mobilenet_v3 import LCNet100
-from kimm.models.mobilenet_v3 import MobileNet100V3Large
-from kimm.models.mobilenet_v3 import MobileNet100V3Small
-from kimm.models.mobilenet_v3 import MobileNet100V3SmallMinimal
-from kimm.utils import make_divisible
-
-
-class MobileNetV3Test(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        [
-            (MobileNet100V3Small.__name__, MobileNet100V3Small),
-            (MobileNet100V3SmallMinimal.__name__, MobileNet100V3SmallMinimal),
-            (MobileNet100V3Large.__name__, MobileNet100V3Large),
-            (LCNet100.__name__, LCNet100),
-        ]
-    )
-    def test_mobilenet_v3_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters(
-        [
-            (MobileNet100V3Small.__name__, MobileNet100V3Small, 1.0),
-            (
-                MobileNet100V3SmallMinimal.__name__,
-                MobileNet100V3SmallMinimal,
-                1.0,
-            ),
-            (MobileNet100V3Large.__name__, MobileNet100V3Large, 1.0),
-        ]
-    )
-    def test_mobilenet_v3_feature_extractor(self, model_class, width):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        if "Small" in model_class.__name__:
-            self.assertEqual(
-                list(y["STEM_S2"].shape),
-                [1, 112, 112, make_divisible(16 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK0_S4"].shape),
-                [1, 56, 56, make_divisible(16 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK1_S8"].shape),
-                [1, 28, 28, make_divisible(24 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK2_S16"].shape),
-                [1, 14, 14, make_divisible(40 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK4_S32"].shape),
-                [1, 7, 7, make_divisible(96 * width)],
-            )
-        else:
-            self.assertEqual(
-                list(y["STEM_S2"].shape),
-                [1, 112, 112, make_divisible(16 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK1_S4"].shape),
-                [1, 56, 56, make_divisible(24 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK2_S8"].shape),
-                [1, 28, 28, make_divisible(40 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK3_S16"].shape),
-                [1, 14, 14, make_divisible(80 * width)],
-            )
-            self.assertEqual(
-                list(y["BLOCK5_S32"].shape),
-                [1, 7, 7, make_divisible(160 * width)],
-            )
-
-    @parameterized.named_parameters([(LCNet100.__name__, LCNet100, 1.0)])
-    def test_lcnet_feature_extractor(self, model_class, width):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(
-            list(y["STEM_S2"].shape),
-            [1, 112, 112, make_divisible(16 * width)],
-        )
-        self.assertEqual(
-            list(y["BLOCK1_S4"].shape),
-            [1, 56, 56, make_divisible(64 * width)],
-        )
-        self.assertEqual(
-            list(y["BLOCK2_S8"].shape),
-            [1, 28, 28, make_divisible(128 * width)],
-        )
-        self.assertEqual(
-            list(y["BLOCK3_S16"].shape),
-            [1, 14, 14, make_divisible(256 * width)],
-        )
-        self.assertEqual(
-            list(y["BLOCK5_S32"].shape),
-            [1, 7, 7, make_divisible(512 * width)],
-        )
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [
-            (MobileNet100V3Small.__name__, MobileNet100V3Small, 224),
-            (
-                MobileNet100V3SmallMinimal.__name__,
-                MobileNet100V3SmallMinimal,
-                224,
-            ),
-            (MobileNet100V3Large.__name__, MobileNet100V3Large, 224),
-            (LCNet100.__name__, LCNet100, 224),
-        ]
-    )
-    def test_mobilenet_v3_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/mobilevit_test.py b/kimm/models/mobilevit_test.py
deleted file mode 100644
index eadb290..0000000
--- a/kimm/models/mobilevit_test.py
+++ /dev/null
@@ -1,68 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.mobilevit import MobileViTS
-from kimm.models.mobilevit import MobileViTXS
-
-
-class MobileViTTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        [(MobileViTS.__name__, MobileViTS), (MobileViTXS.__name__, MobileViTXS)]
-    )
-    def test_mobilevit_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 256, 256, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters(
-        [(MobileViTS.__name__, MobileViTS), (MobileViTXS.__name__, MobileViTXS)]
-    )
-    def test_mobilevit_feature_extractor(self, model_class):
-        x = random.uniform([1, 256, 256, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        if "MobileViTS" in model_class.__name__:
-            self.assertEqual(list(y["STEM_S2"].shape), [1, 128, 128, 16])
-            self.assertEqual(list(y["BLOCK1_S4"].shape), [1, 64, 64, 64])
-            self.assertEqual(list(y["BLOCK2_S8"].shape), [1, 32, 32, 96])
-            self.assertEqual(list(y["BLOCK3_S16"].shape), [1, 16, 16, 128])
-            self.assertEqual(list(y["BLOCK4_S32"].shape), [1, 8, 8, 160])
-        elif "MobileViTXS" in model_class.__name__:
-            self.assertEqual(list(y["STEM_S2"].shape), [1, 128, 128, 16])
-            self.assertEqual(list(y["BLOCK1_S4"].shape), [1, 64, 64, 48])
-            self.assertEqual(list(y["BLOCK2_S8"].shape), [1, 32, 32, 64])
-            self.assertEqual(list(y["BLOCK3_S16"].shape), [1, 16, 16, 80])
-            self.assertEqual(list(y["BLOCK4_S32"].shape), [1, 8, 8, 96])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [
-            (MobileViTS.__name__, MobileViTS, 256),
-            (MobileViTXS.__name__, MobileViTXS, 256),
-        ]
-    )
-    def test_mobilevit_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/models_test.py b/kimm/models/models_test.py
new file mode 100644
index 0000000..d89f4d7
--- /dev/null
+++ b/kimm/models/models_test.py
@@ -0,0 +1,364 @@
+import pytest
+from absl.testing import parameterized
+from keras import models
+from keras import random
+from keras.src import testing
+
+from kimm import models as kimm_models
+from kimm.utils import make_divisible
+
+# name, class, default_size, features (name, shape)
+MODEL_CONFIGS = [
+    # convmixer
+    (
+        kimm_models.ConvMixer736D32.__name__,
+        kimm_models.ConvMixer736D32,
+        224,
+        [
+            ("STEM", [1, 32, 32, 768]),
+            *((f"BLOCK{i}", [1, 32, 32, 768]) for i in range(32)),
+        ],
+    ),
+    # densenet
+    (
+        kimm_models.DenseNet121.__name__,
+        kimm_models.DenseNet121,
+        224,
+        [
+            ("STEM_S4", [1, 56, 56, 64]),
+            ("BLOCK0_S8", [1, 28, 28, 128]),
+            ("BLOCK1_S16", [1, 14, 14, 256]),
+            ("BLOCK2_S32", [1, 7, 7, 512]),
+            ("BLOCK3_S32", [1, 7, 7, 1024]),
+        ],
+    ),
+    # efficientnet
+    (
+        kimm_models.EfficientNetB2.__name__,
+        kimm_models.EfficientNetB2,
+        260,
+        [
+            ("STEM_S2", [1, 130, 130, make_divisible(32 * 1.1)]),
+            ("BLOCK1_S4", [1, 65, 65, make_divisible(24 * 1.1)]),
+            ("BLOCK2_S8", [1, 33, 33, make_divisible(40 * 1.1)]),
+            ("BLOCK3_S16", [1, 17, 17, make_divisible(80 * 1.1)]),
+            ("BLOCK5_S32", [1, 9, 9, make_divisible(192 * 1.1)]),
+        ],
+    ),
+    (
+        kimm_models.EfficientNetLiteB2.__name__,
+        kimm_models.EfficientNetLiteB2,
+        260,
+        [
+            ("STEM_S2", [1, 130, 130, make_divisible(32 * 1.1)]),
+            ("BLOCK1_S4", [1, 65, 65, make_divisible(24 * 1.1)]),
+            ("BLOCK2_S8", [1, 33, 33, make_divisible(40 * 1.1)]),
+            ("BLOCK3_S16", [1, 17, 17, make_divisible(80 * 1.1)]),
+            ("BLOCK5_S32", [1, 9, 9, make_divisible(192 * 1.1)]),
+        ],
+    ),
+    (
+        kimm_models.EfficientNetV2S.__name__,
+        kimm_models.EfficientNetV2S,
+        300,
+        [
+            ("STEM_S2", [1, 150, 150, make_divisible(24 * 1.0)]),
+            ("BLOCK1_S4", [1, 75, 75, make_divisible(48 * 1.0)]),
+            ("BLOCK2_S8", [1, 38, 38, make_divisible(64 * 1.0)]),
+            ("BLOCK3_S16", [1, 19, 19, make_divisible(128 * 1.0)]),
+            ("BLOCK5_S32", [1, 10, 10, make_divisible(256 * 1.0)]),
+        ],
+    ),
+    (
+        kimm_models.EfficientNetV2B0.__name__,
+        kimm_models.EfficientNetV2B0,
+        192,
+        [
+            ("STEM_S2", [1, 96, 96, make_divisible(32 * 1.0)]),
+            ("BLOCK1_S4", [1, 48, 48, make_divisible(32 * 1.0)]),
+            ("BLOCK2_S8", [1, 24, 24, make_divisible(48 * 1.0)]),
+            ("BLOCK3_S16", [1, 12, 12, make_divisible(96 * 1.0)]),
+            ("BLOCK5_S32", [1, 6, 6, make_divisible(192 * 1.0)]),
+        ],
+    ),
+    (
+        kimm_models.TinyNetE.__name__,
+        kimm_models.TinyNetE,
+        106,
+        [
+            ("STEM_S2", [1, 53, 53, 32]),
+            ("BLOCK1_S4", [1, 27, 27, make_divisible(24 * 0.51)]),
+            ("BLOCK2_S8", [1, 14, 14, make_divisible(40 * 0.51)]),
+            ("BLOCK3_S16", [1, 7, 7, make_divisible(80 * 0.51)]),
+            ("BLOCK5_S32", [1, 4, 4, make_divisible(192 * 0.51)]),
+        ],
+    ),
+    # ghostnet
+    (
+        kimm_models.GhostNet100.__name__,
+        kimm_models.GhostNet100,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, 16]),
+            ("BLOCK1_S4", [1, 56, 56, 24]),
+            ("BLOCK3_S8", [1, 28, 28, 40]),
+            ("BLOCK5_S16", [1, 14, 14, 80]),
+            ("BLOCK7_S32", [1, 7, 7, 160]),
+        ],
+    ),
+    (
+        kimm_models.GhostNet100V2.__name__,
+        kimm_models.GhostNet100V2,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, 16]),
+            ("BLOCK1_S4", [1, 56, 56, 24]),
+            ("BLOCK3_S8", [1, 28, 28, 40]),
+            ("BLOCK5_S16", [1, 14, 14, 80]),
+            ("BLOCK7_S32", [1, 7, 7, 160]),
+        ],
+    ),
+    # inception_v3
+    (
+        kimm_models.InceptionV3.__name__,
+        kimm_models.InceptionV3,
+        299,
+        [
+            ("STEM_S2", [1, 147, 147, 64]),
+            ("BLOCK0_S4", [1, 71, 71, 192]),
+            ("BLOCK1_S8", [1, 35, 35, 288]),
+            ("BLOCK2_S16", [1, 17, 17, 768]),
+            ("BLOCK3_S32", [1, 8, 8, 2048]),
+        ],
+    ),
+    # mobilenet_v2
+    (
+        kimm_models.MobileNet050V2.__name__,
+        kimm_models.MobileNet050V2,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, make_divisible(32 * 0.5)]),
+            ("BLOCK1_S4", [1, 56, 56, make_divisible(24 * 0.5)]),
+            ("BLOCK2_S8", [1, 28, 28, make_divisible(32 * 0.5)]),
+            ("BLOCK3_S16", [1, 14, 14, make_divisible(64 * 0.5)]),
+            ("BLOCK5_S32", [1, 7, 7, make_divisible(160 * 0.5)]),
+        ],
+    ),
+    (
+        kimm_models.MobileNet100V2.__name__,
+        kimm_models.MobileNet100V2,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, make_divisible(32 * 1.0)]),
+            ("BLOCK1_S4", [1, 56, 56, make_divisible(24 * 1.0)]),
+            ("BLOCK2_S8", [1, 28, 28, make_divisible(32 * 1.0)]),
+            ("BLOCK3_S16", [1, 14, 14, make_divisible(64 * 1.0)]),
+            ("BLOCK5_S32", [1, 7, 7, make_divisible(160 * 1.0)]),
+        ],
+    ),
+    # mobilenet_v3
+    (
+        kimm_models.LCNet100.__name__,
+        kimm_models.LCNet100,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, make_divisible(16 * 1.0)]),
+            ("BLOCK1_S4", [1, 56, 56, make_divisible(64 * 1.0)]),
+            ("BLOCK2_S8", [1, 28, 28, make_divisible(128 * 1.0)]),
+            ("BLOCK3_S16", [1, 14, 14, make_divisible(256 * 1.0)]),
+            ("BLOCK5_S32", [1, 7, 7, make_divisible(512 * 1.0)]),
+        ],
+    ),
+    (
+        kimm_models.MobileNet100V3Large.__name__,
+        kimm_models.MobileNet100V3Large,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, make_divisible(16 * 1.0)]),
+            ("BLOCK1_S4", [1, 56, 56, make_divisible(24 * 1.0)]),
+            ("BLOCK2_S8", [1, 28, 28, make_divisible(40 * 1.0)]),
+            ("BLOCK3_S16", [1, 14, 14, make_divisible(80 * 1.0)]),
+            ("BLOCK5_S32", [1, 7, 7, make_divisible(160 * 1.0)]),
+        ],
+    ),
+    (
+        kimm_models.MobileNet100V3Small.__name__,
+        kimm_models.MobileNet100V3Small,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, make_divisible(16 * 1.0)]),
+            ("BLOCK0_S4", [1, 56, 56, make_divisible(16 * 1.0)]),
+            ("BLOCK1_S8", [1, 28, 28, make_divisible(24 * 1.0)]),
+            ("BLOCK2_S16", [1, 14, 14, make_divisible(40 * 1.0)]),
+            ("BLOCK4_S32", [1, 7, 7, make_divisible(96 * 1.0)]),
+        ],
+    ),
+    (
+        kimm_models.MobileNet100V3SmallMinimal.__name__,
+        kimm_models.MobileNet100V3SmallMinimal,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, make_divisible(16 * 1.0)]),
+            ("BLOCK0_S4", [1, 56, 56, make_divisible(16 * 1.0)]),
+            ("BLOCK1_S8", [1, 28, 28, make_divisible(24 * 1.0)]),
+            ("BLOCK2_S16", [1, 14, 14, make_divisible(40 * 1.0)]),
+            ("BLOCK4_S32", [1, 7, 7, make_divisible(96 * 1.0)]),
+        ],
+    ),
+    # mobilevit
+    (
+        kimm_models.MobileViTS.__name__,
+        kimm_models.MobileViTS,
+        256,
+        [
+            ("STEM_S2", [1, 128, 128, 16]),
+            ("BLOCK1_S4", [1, 64, 64, 64]),
+            ("BLOCK2_S8", [1, 32, 32, 96]),
+            ("BLOCK3_S16", [1, 16, 16, 128]),
+            ("BLOCK4_S32", [1, 8, 8, 160]),
+        ],
+    ),
+    (
+        kimm_models.MobileViTXS.__name__,
+        kimm_models.MobileViTXS,
+        256,
+        [
+            ("STEM_S2", [1, 128, 128, 16]),
+            ("BLOCK1_S4", [1, 64, 64, 48]),
+            ("BLOCK2_S8", [1, 32, 32, 64]),
+            ("BLOCK3_S16", [1, 16, 16, 80]),
+            ("BLOCK4_S32", [1, 8, 8, 96]),
+        ],
+    ),
+    # regnet
+    (
+        kimm_models.RegNetX002.__name__,
+        kimm_models.RegNetX002,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, 32]),
+            ("BLOCK0_S4", [1, 56, 56, 24]),
+            ("BLOCK1_S8", [1, 28, 28, 56]),
+            ("BLOCK2_S16", [1, 14, 14, 152]),
+            ("BLOCK3_S32", [1, 7, 7, 368]),
+        ],
+    ),
+    (
+        kimm_models.RegNetY002.__name__,
+        kimm_models.RegNetY002,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, 32]),
+            ("BLOCK0_S4", [1, 56, 56, 24]),
+            ("BLOCK1_S8", [1, 28, 28, 56]),
+            ("BLOCK2_S16", [1, 14, 14, 152]),
+            ("BLOCK3_S32", [1, 7, 7, 368]),
+        ],
+    ),
+    # resnet
+    (
+        kimm_models.ResNet18.__name__,
+        kimm_models.ResNet18,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, 64]),
+            ("BLOCK0_S4", [1, 56, 56, 64]),
+            ("BLOCK1_S8", [1, 28, 28, 128]),
+            ("BLOCK2_S16", [1, 14, 14, 256]),
+            ("BLOCK3_S32", [1, 7, 7, 512]),
+        ],
+    ),
+    (
+        kimm_models.ResNet50.__name__,
+        kimm_models.ResNet50,
+        224,
+        [
+            ("STEM_S2", [1, 112, 112, 64]),
+            ("BLOCK0_S4", [1, 56, 56, 64 * 4]),
+            ("BLOCK1_S8", [1, 28, 28, 128 * 4]),
+            ("BLOCK2_S16", [1, 14, 14, 256 * 4]),
+            ("BLOCK3_S32", [1, 7, 7, 512 * 4]),
+        ],
+    ),
+    # vgg
+    (
+        kimm_models.VGG11.__name__,
+        kimm_models.VGG11,
+        224,
+        [
+            ("BLOCK0_S1", [1, 224, 224, 64]),
+            ("BLOCK1_S2", [1, 112, 112, 128]),
+            ("BLOCK2_S4", [1, 56, 56, 256]),
+            ("BLOCK3_S8", [1, 28, 28, 512]),
+            ("BLOCK4_S16", [1, 14, 14, 512]),
+            ("BLOCK5_S32", [1, 7, 7, 512]),
+        ],
+    ),
+    # vision_transformer
+    (
+        kimm_models.VisionTransformerTiny16.__name__,
+        kimm_models.VisionTransformerTiny16,
+        384,
+        [*((f"BLOCK{i}", [1, 577, 192]) for i in range(5))],
+    ),
+    (
+        kimm_models.VisionTransformerTiny32.__name__,
+        kimm_models.VisionTransformerTiny32,
+        384,
+        [*((f"BLOCK{i}", [1, 145, 192]) for i in range(5))],
+    ),
+    # xception
+    (
+        kimm_models.Xception.__name__,
+        kimm_models.Xception,
+        299,
+        [
+            ("STEM_S2", [1, 147, 147, 64]),
+            ("BLOCK0_S4", [1, 74, 74, 128]),
+            ("BLOCK1_S8", [1, 37, 37, 256]),
+            ("BLOCK2_S16", [1, 19, 19, 728]),
+            ("BLOCK3_S32", [1, 10, 10, 2048]),
+        ],
+    ),
+]
+
+
+class ConvMixerTest(testing.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(MODEL_CONFIGS)
+    def test_model_base(self, model_class, image_size, features):
+        # TODO: test the correctness of the real image
+        x = random.uniform([1, image_size, image_size, 3]) * 255.0
+        model = model_class()
+
+        y = model(x, training=False)
+
+        self.assertEqual(y.shape, (1, 1000))
+
+    @parameterized.named_parameters(MODEL_CONFIGS)
+    def test_model_feature_extractor(self, model_class, image_size, features):
+        x = random.uniform([1, image_size, image_size, 3]) * 255.0
+        model = model_class(feature_extractor=True)
+
+        y = model(x, training=False)
+
+        self.assertIsInstance(y, dict)
+        self.assertContainsSubset(
+            model_class.available_feature_keys(), list(y.keys())
+        )
+        for feature_info in features:
+            name, shape = feature_info
+            self.assertEqual(list(y[name].shape), shape)
+
+    @pytest.mark.serialization
+    @parameterized.named_parameters(MODEL_CONFIGS)
+    def test_model_serialization(self, model_class, image_size, features):
+        x = random.uniform([1, image_size, image_size, 3]) * 255.0
+        temp_dir = self.get_temp_dir()
+        model1 = model_class()
+
+        y1 = model1(x, training=False)
+        model1.save(temp_dir + "/model.keras")
+        model2 = models.load_model(temp_dir + "/model.keras")
+        y2 = model2(x, training=False)
+
+        self.assertAllClose(y1, y2)
diff --git a/kimm/models/regnet_test.py b/kimm/models/regnet_test.py
deleted file mode 100644
index 906586b..0000000
--- a/kimm/models/regnet_test.py
+++ /dev/null
@@ -1,61 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.regnet import RegNetX002
-from kimm.models.regnet import RegNetY002
-
-
-class RegNetTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        [(RegNetX002.__name__, RegNetX002), (RegNetY002.__name__, RegNetY002)]
-    )
-    def test_regnet_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters(
-        [(RegNetX002.__name__, RegNetX002), (RegNetY002.__name__, RegNetY002)]
-    )
-    def test_regnet_feature_extractor(self, model_class):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(list(y["STEM_S2"].shape), [1, 112, 112, 32])
-        self.assertEqual(list(y["BLOCK0_S4"].shape), [1, 56, 56, 24])
-        self.assertEqual(list(y["BLOCK1_S8"].shape), [1, 28, 28, 56])
-        self.assertEqual(list(y["BLOCK2_S16"].shape), [1, 14, 14, 152])
-        self.assertEqual(list(y["BLOCK3_S32"].shape), [1, 7, 7, 368])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [
-            (RegNetX002.__name__, RegNetX002, 224),
-            (RegNetY002.__name__, RegNetY002, 224),
-        ]
-    )
-    def test_regnet_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/resnet_test.py b/kimm/models/resnet_test.py
deleted file mode 100644
index ec06921..0000000
--- a/kimm/models/resnet_test.py
+++ /dev/null
@@ -1,66 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.resnet import ResNet18
-from kimm.models.resnet import ResNet50
-
-
-class ResNetTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        [(ResNet18.__name__, ResNet18), (ResNet50.__name__, ResNet50)]
-    )
-    def test_resnet_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters(
-        [(ResNet18.__name__, ResNet18, 1), (ResNet50.__name__, ResNet50, 4)]
-    )
-    def test_resnet_feature_extractor(self, model_class, expansion):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(list(y["STEM_S2"].shape), [1, 112, 112, 64])
-        self.assertEqual(
-            list(y["BLOCK0_S4"].shape), [1, 56, 56, 64 * expansion]
-        )
-        self.assertEqual(
-            list(y["BLOCK1_S8"].shape), [1, 28, 28, 128 * expansion]
-        )
-        self.assertEqual(
-            list(y["BLOCK2_S16"].shape), [1, 14, 14, 256 * expansion]
-        )
-        self.assertEqual(
-            list(y["BLOCK3_S32"].shape), [1, 7, 7, 512 * expansion]
-        )
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [(ResNet18.__name__, ResNet18, 224), (ResNet50.__name__, ResNet50, 224)]
-    )
-    def test_resnet_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/vgg_test.py b/kimm/models/vgg_test.py
deleted file mode 100644
index cafa67d..0000000
--- a/kimm/models/vgg_test.py
+++ /dev/null
@@ -1,52 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.vgg import VGG11
-
-
-class VGGTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters([(VGG11.__name__, VGG11)])
-    def test_densenet_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters([(VGG11.__name__, VGG11)])
-    def test_densenet_feature_extractor(self, model_class):
-        x = random.uniform([1, 224, 224, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(list(y["BLOCK0_S1"].shape), [1, 224, 224, 64])
-        self.assertEqual(list(y["BLOCK1_S2"].shape), [1, 112, 112, 128])
-        self.assertEqual(list(y["BLOCK2_S4"].shape), [1, 56, 56, 256])
-        self.assertEqual(list(y["BLOCK3_S8"].shape), [1, 28, 28, 512])
-        self.assertEqual(list(y["BLOCK4_S16"].shape), [1, 14, 14, 512])
-        self.assertEqual(list(y["BLOCK5_S32"].shape), [1, 7, 7, 512])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters([(VGG11.__name__, VGG11, 224)])
-    def test_densenet_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/vision_transformer_test.py b/kimm/models/vision_transformer_test.py
deleted file mode 100644
index bfd87fe..0000000
--- a/kimm/models/vision_transformer_test.py
+++ /dev/null
@@ -1,72 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.vision_transformer import VisionTransformerTiny16
-from kimm.models.vision_transformer import VisionTransformerTiny32
-
-
-class VisionTransformerTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        [
-            (VisionTransformerTiny16.__name__, VisionTransformerTiny16),
-            (VisionTransformerTiny32.__name__, VisionTransformerTiny32),
-        ]
-    )
-    def test_vision_transformer_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 384, 384, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters(
-        [
-            (VisionTransformerTiny16.__name__, VisionTransformerTiny16, 16),
-            (VisionTransformerTiny32.__name__, VisionTransformerTiny32, 32),
-        ]
-    )
-    def test_vision_transformer_feature_extractor(
-        self, model_class, patch_size
-    ):
-        x = random.uniform([1, 384, 384, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        if patch_size == 16:
-            self.assertEqual(list(y["BLOCK0"].shape), [1, 577, 192])
-        elif patch_size == 32:
-            self.assertEqual(list(y["BLOCK0"].shape), [1, 145, 192])
-        if patch_size == 16:
-            self.assertEqual(list(y["BLOCK5"].shape), [1, 577, 192])
-        elif patch_size == 32:
-            self.assertEqual(list(y["BLOCK5"].shape), [1, 145, 192])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters(
-        [
-            (VisionTransformerTiny16.__name__, VisionTransformerTiny16, 384),
-            (VisionTransformerTiny32.__name__, VisionTransformerTiny32, 384),
-        ]
-    )
-    def test_vit_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)
diff --git a/kimm/models/xception_test.py b/kimm/models/xception_test.py
deleted file mode 100644
index d88e403..0000000
--- a/kimm/models/xception_test.py
+++ /dev/null
@@ -1,51 +0,0 @@
-import pytest
-from absl.testing import parameterized
-from keras import models
-from keras import random
-from keras.src import testing
-
-from kimm.models.xception import Xception
-
-
-class XceptionTest(testing.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters([(Xception.__name__, Xception)])
-    def test_densenet_base(self, model_class):
-        # TODO: test the correctness of the real image
-        x = random.uniform([1, 299, 299, 3]) * 255.0
-        model = model_class()
-
-        y = model(x, training=False)
-
-        self.assertEqual(y.shape, (1, 1000))
-
-    @parameterized.named_parameters([(Xception.__name__, Xception)])
-    def test_densenet_feature_extractor(self, model_class):
-        x = random.uniform([1, 299, 299, 3]) * 255.0
-        model = model_class(feature_extractor=True)
-
-        y = model(x, training=False)
-
-        self.assertIsInstance(y, dict)
-        self.assertContainsSubset(
-            model_class.available_feature_keys(),
-            list(y.keys()),
-        )
-        self.assertEqual(list(y["STEM_S2"].shape), [1, 147, 147, 64])
-        self.assertEqual(list(y["BLOCK0_S4"].shape), [1, 74, 74, 128])
-        self.assertEqual(list(y["BLOCK1_S8"].shape), [1, 37, 37, 256])
-        self.assertEqual(list(y["BLOCK2_S16"].shape), [1, 19, 19, 728])
-        self.assertEqual(list(y["BLOCK3_S32"].shape), [1, 10, 10, 2048])
-
-    @pytest.mark.serialization
-    @parameterized.named_parameters([(Xception.__name__, Xception, 299)])
-    def test_densenet_serialization(self, model_class, image_size):
-        x = random.uniform([1, image_size, image_size, 3]) * 255.0
-        temp_dir = self.get_temp_dir()
-        model1 = model_class()
-        y1 = model1(x, training=False)
-        model1.save(temp_dir + "/model.keras")
-
-        model2 = models.load_model(temp_dir + "/model.keras")
-        y2 = model2(x, training=False)
-
-        self.assertAllClose(y1, y2)