Refactor model definition (#42)

* Refactor

* Update version

* Untrack `GlobalResponseNormalization`

* Add warning if no available pretrained weights

* Allow UserWarning

Author: james77777778

.github/workflows/actions.yml

@@ -12,6 +12,17 @@ permissions:
   contents: read
 
 jobs:
+  format:
+    name: Check the code format
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.9'
+      - uses: pre-commit/action@v3.0.1
+
   build:
     strategy:
       fail-fast: false
@@ -54,29 +65,3 @@ jobs:
           files: coverage.xml
           flags: kimm,kimm-${{ matrix.backend }}
           fail_ci_if_error: false
-
-  format:
-    name: Check the code format
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python 3.9
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.9'
-      - name: Get pip cache dir
-        id: Pip-cache
-        run: |
-          python -m pip install --upgrade pip setuptools
-          echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
-      - name: pip cache
-        uses: actions/cache@v4
-        with:
-          path: ${{ steps.pip-cache.outputs.dir }}
-          key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}-${{ hashFiles('requirements.txt') }}
-      - name: Install dependencies
-        run: |
-          pip install -r requirements.txt --progress-bar off --upgrade
-          pip install -e ".[tests]" --progress-bar off --upgrade
-      - name: Lint
-        run: bash shell/lint.sh

.pre-commit-config.yaml

@@ -0,0 +1,33 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.6.0
+  hooks:
+  - id: check-ast
+  - id: check-merge-conflict
+  - id: check-toml
+  - id: check-yaml
+  - id: end-of-file-fixer
+    files: \.py$
+  - id: debug-statements
+    files: \.py$
+  - id: trailing-whitespace
+    files: \.py$
+
+- repo: https://github.com/pycqa/isort
+  rev: 5.13.2
+  hooks:
+  - id: isort
+    name: isort (python)
+
+- repo: https://github.com/psf/black-pre-commit-mirror
+  rev: 24.4.2
+  hooks:
+  - id: black
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.4.4
+  hooks:
+  - id: ruff
+    args:
+      - --fix
+  - id: ruff-format

kimm/__init__.py

@@ -3,4 +3,4 @@
 from kimm import utils
 from kimm.utils.model_registry import list_models
 
-__version__ = "0.1.8"
+__version__ = "0.2.0"

kimm/models/convmixer.py

@@ -134,20 +134,16 @@ def fix_config(self, config):
         return config
 
 
-"""
-Model Definition
-"""
+# Model Definition
 
 
-class ConvMixer736D32(ConvMixer):
-    available_feature_keys = ["STEM", *[f"BLOCK{i}" for i in range(32)]]
-    available_weights = [
-        (
-            "imagenet",
-            ConvMixer.default_origin,
-            "convmixer736d32_convmixer_768_32.in1k.keras",
-        )
-    ]
+class ConvMixerVariant(ConvMixer):
+    # Parameters
+    depth = None
+    hidden_channels = None
+    patch_size = None
+    kernel_size = None
+    activation = None
 
     def __init__(
         self,
@@ -160,16 +156,21 @@ def __init__(
         classes: int = 1000,
         classifier_activation: str = "softmax",
         weights: typing.Optional[str] = "imagenet",
-        name: str = "ConvMixer736D32",
+        name: typing.Optional[str] = None,
         **kwargs,
     ):
+        if type(self) is ConvMixerVariant:
+            raise NotImplementedError(
+                f"Cannot instantiate base class: {self.__class__.__name__}. "
+                "You should use its subclasses."
+            )
         kwargs = self.fix_config(kwargs)
         super().__init__(
-            32,
-            768,
-            7,
-            7,
-            "relu",
+            depth=self.depth,
+            hidden_channels=self.hidden_channels,
+            patch_size=self.patch_size,
+            kernel_size=self.kernel_size,
+            activation=self.activation,
             input_tensor=input_tensor,
             input_shape=input_shape,
             include_preprocessing=include_preprocessing,
@@ -179,12 +180,30 @@ def __init__(
             classes=classes,
             classifier_activation=classifier_activation,
             weights=weights,
-            name=name,
+            name=name or str(self.__class__.__name__),
             **kwargs,
         )
 
 
-class ConvMixer1024D20(ConvMixer):
+class ConvMixer736D32(ConvMixerVariant):
+    available_feature_keys = ["STEM", *[f"BLOCK{i}" for i in range(32)]]
+    available_weights = [
+        (
+            "imagenet",
+            ConvMixer.default_origin,
+            "convmixer736d32_convmixer_768_32.in1k.keras",
+        )
+    ]
+
+    # Parameters
+    depth = 32
+    hidden_channels = 768
+    patch_size = 7
+    kernel_size = 7
+    activation = "relu"
+
+
+class ConvMixer1024D20(ConvMixerVariant):
     available_feature_keys = ["STEM", *[f"BLOCK{i}" for i in range(20)]]
     available_weights = [
         (
@@ -194,42 +213,15 @@ 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] = "imagenet",
-        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,
-        )
+    # Parameters
+    depth = 20
+    hidden_channels = 1024
+    patch_size = 14
+    kernel_size = 9
+    activation = "gelu"
 
 
-class ConvMixer1536D20(ConvMixer):
+class ConvMixer1536D20(ConvMixerVariant):
     available_feature_keys = ["STEM", *[f"BLOCK{i}" for i in range(20)]]
     available_weights = [
         (
@@ -239,39 +231,12 @@ 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] = "imagenet",
-        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,
-        )
+    # Parameters
+    depth = 20
+    hidden_channels = 1536
+    patch_size = 7
+    kernel_size = 9
+    activation = "gelu"
 
 
 add_model_to_registry(ConvMixer736D32, "imagenet")