Refactor APIs using namex (#43)

* Refactor `kimm.blocks`

* Update

* Export model apis

* Update apis

* Fix version

* Update CI

* Update CI

* Update `release.yml`

* Update tests

* Speed up tests

* Update keras version

* Update keras version

.github/workflows/actions.yml

@@ -21,35 +21,57 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: '3.9'
-      - uses: pre-commit/action@v3.0.1
+      - name: Lint
+        uses: pre-commit/action@v3.0.1
+      - name: Get pip cache dir
+        id: pip-cache
+        run: |
+          python -m pip install --upgrade pip setuptools
+          echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
+      - name: Cache pip
+        uses: actions/cache@v4
+        with:
+          path: ${{ steps.pip-cache.outputs.dir }}
+          key: ${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }}-${{ 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: Check for API changes
+        run: |
+          bash shell/api_gen.sh
+          git status
+          clean=$(git status | grep "nothing to commit")
+          if [ -z "$clean" ]; then
+            echo "Please run shell/api_gen.sh to generate API."
+            exit 1
+          fi
 
   build:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.9]
         backend: [tensorflow, jax, torch, numpy]
     name: Run tests
     runs-on: ubuntu-latest
     env:
-      PYTHON: ${{ matrix.python-version }}
       KERAS_BACKEND: ${{ matrix.backend }}
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python
+      - name: Set up Python 3.9
         uses: actions/setup-python@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          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
+      - name: Cache pip
         uses: actions/cache@v4
         with:
           path: ${{ steps.pip-cache.outputs.dir }}
-          key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}-${{ hashFiles('requirements.txt') }}
+          key: ${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('requirements.txt') }}
       - name: Install dependencies
         run: |
           pip install -r requirements.txt --progress-bar off --upgrade

.github/workflows/release.yml

@@ -23,11 +23,11 @@ jobs:
         run: |
           python -m pip install --upgrade pip setuptools
           echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
-      - name: Pip cache
+      - name: Cache pip
         uses: actions/cache@v4
         with:
           path: ${{ steps.pip-cache.outputs.dir }}
-          key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}-${{ hashFiles('requirements.txt') }}
+          key: ${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('requirements.txt') }}
       - name: Install dependencies
         run: |
           pip install -r requirements.txt --progress-bar off --upgrade

.pre-commit-config.yaml

@@ -17,7 +17,6 @@ repos:
   rev: 5.13.2
   hooks:
   - id: isort
-    name: isort (python)
 
 - repo: https://github.com/psf/black-pre-commit-mirror
   rev: 24.4.2

README.md

@@ -4,7 +4,7 @@
 <div align="center">
 <img width="50%" src="https://github.com/james77777778/kimm/assets/20734616/b21db8f2-307b-4791-b93d-e913e45fb238" alt="KIMM">
 
-[![Keras](https://img.shields.io/badge/keras-v3.0.4+-success.svg)](https://github.com/keras-team/keras)
+[![Keras](https://img.shields.io/badge/keras-v3.3.0+-success.svg)](https://github.com/keras-team/keras)
 [![PyPI](https://img.shields.io/pypi/v/kimm)](https://pypi.org/project/kimm/)
 [![Contributions Welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/james77777778/kimm/issues)
 [![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/james77777778/keras-image-models/actions.yml?label=tests)](https://github.com/james77777778/keras-image-models/actions/workflows/actions.yml?query=branch%3Amain++)

api_gen.py

@@ -0,0 +1,13 @@
+import namex
+
+from kimm._src.version import __version__
+
+namex.generate_api_files(package="kimm", code_directory="_src")
+
+# Add version string
+
+with open("kimm/__init__.py", "r") as f:
+    contents = f.read()
+with open("kimm/__init__.py", "w") as f:
+    contents += f'__version__ = "{__version__}"\n'
+    f.write(contents)

kimm/__init__.py

@@ -1,6 +1,16 @@
+"""DO NOT EDIT.
+
+This file was autogenerated. Do not edit it by hand,
+since your modifications would be overwritten.
+"""
+
+from kimm import blocks
 from kimm import export
+from kimm import layers
 from kimm import models
+from kimm import timm_utils
 from kimm import utils
-from kimm.utils.model_registry import list_models
+from kimm._src.utils.model_registry import list_models
+from kimm._src.version import version
 
 __version__ = "0.2.0"

kimm/blocks/base_block.py → kimm/_src/blocks/conv2d.py

@@ -3,25 +3,10 @@
 from keras import backend
 from keras import layers
 
-from kimm.utils import make_divisible
-
-
-def apply_activation(
-    inputs, activation: typing.Optional[str] = None, name: str = "activation"
-):
-    x = inputs
-    if activation is not None:
-        if isinstance(activation, str):
-            x = layers.Activation(activation, name=name)(x)
-        elif isinstance(activation, layers.Layer):
-            x = activation(x)
-        else:
-            NotImplementedError(
-                f"Unsupported activation type: {type(activation)}"
-            )
-    return x
+from kimm._src.kimm_export import kimm_export
 
 
+@kimm_export(parent_path=["kimm.blocks"])
 def apply_conv2d_block(
     inputs,
     filters: typing.Optional[int] = None,
@@ -83,45 +68,8 @@ def apply_conv2d_block(
         momentum=bn_momentum,
         epsilon=bn_epsilon,
     )(x)
-    x = apply_activation(x, activation, name=name)
+    if activation is not None:
+        x = layers.Activation(activation, name=name)(x)
     if has_skip:
         x = layers.Add()([x, inputs])
     return x
-
-
-def apply_se_block(
-    inputs,
-    se_ratio: float = 0.25,
-    activation: typing.Optional[str] = "relu",
-    gate_activation: typing.Optional[str] = "sigmoid",
-    make_divisible_number: typing.Optional[int] = None,
-    se_input_channels: typing.Optional[int] = None,
-    name: str = "se_block",
-):
-    channels_axis = -1 if backend.image_data_format() == "channels_last" else -3
-    input_channels = inputs.shape[channels_axis]
-    if se_input_channels is None:
-        se_input_channels = input_channels
-    if make_divisible_number is None:
-        se_channels = round(se_input_channels * se_ratio)
-    else:
-        se_channels = make_divisible(
-            se_input_channels * se_ratio, make_divisible_number
-        )
-
-    x = inputs
-    x = layers.GlobalAveragePooling2D(
-        data_format=backend.image_data_format(),
-        keepdims=True,
-        name=f"{name}_mean",
-    )(x)
-    x = layers.Conv2D(
-        se_channels, 1, use_bias=True, name=f"{name}_conv_reduce"
-    )(x)
-    x = apply_activation(x, activation, name=f"{name}_act1")
-    x = layers.Conv2D(
-        input_channels, 1, use_bias=True, name=f"{name}_conv_expand"
-    )(x)
-    x = apply_activation(x, gate_activation, name=f"{name}_gate")
-    x = layers.Multiply(name=name)([inputs, x])
-    return x

kimm/blocks/depthwise_separation_block.py → kimm/_src/blocks/depthwise_separation.py

@@ -3,10 +3,12 @@
 from keras import backend
 from keras import layers
 
-from kimm.blocks.base_block import apply_conv2d_block
-from kimm.blocks.base_block import apply_se_block
+from kimm._src.blocks.conv2d import apply_conv2d_block
+from kimm._src.blocks.squeeze_and_excitation import apply_se_block
+from kimm._src.kimm_export import kimm_export
 
 
+@kimm_export(parent_path=["kimm.blocks"])
 def apply_depthwise_separation_block(
     inputs,
     output_channels: int,

kimm/blocks/inverted_residual_block.py → kimm/_src/blocks/inverted_residual.py

@@ -3,11 +3,13 @@
 from keras import backend
 from keras import layers
 
-from kimm.blocks.base_block import apply_conv2d_block
-from kimm.blocks.base_block import apply_se_block
-from kimm.utils import make_divisible
+from kimm._src.blocks.conv2d import apply_conv2d_block
+from kimm._src.blocks.squeeze_and_excitation import apply_se_block
+from kimm._src.kimm_export import kimm_export
+from kimm._src.utils.make_divisble import make_divisible
 
 
+@kimm_export(parent_path=["kimm.blocks"])
 def apply_inverted_residual_block(
     inputs,
     output_channels: int,

kimm/_src/blocks/squeeze_and_excitation.py

@@ -0,0 +1,48 @@
+import typing
+
+from keras import backend
+from keras import layers
+
+from kimm._src.kimm_export import kimm_export
+from kimm._src.utils.make_divisble import make_divisible
+
+
+@kimm_export(parent_path=["kimm.blocks"])
+def apply_se_block(
+    inputs,
+    se_ratio: float = 0.25,
+    activation: typing.Optional[str] = "relu",
+    gate_activation: typing.Optional[str] = "sigmoid",
+    make_divisible_number: typing.Optional[int] = None,
+    se_input_channels: typing.Optional[int] = None,
+    name: str = "se_block",
+):
+    channels_axis = -1 if backend.image_data_format() == "channels_last" else -3
+    input_channels = inputs.shape[channels_axis]
+    if se_input_channels is None:
+        se_input_channels = input_channels
+    if make_divisible_number is None:
+        se_channels = round(se_input_channels * se_ratio)
+    else:
+        se_channels = make_divisible(
+            se_input_channels * se_ratio, make_divisible_number
+        )
+
+    x = inputs
+    x = layers.GlobalAveragePooling2D(
+        data_format=backend.image_data_format(),
+        keepdims=True,
+        name=f"{name}_mean",
+    )(x)
+    x = layers.Conv2D(
+        se_channels, 1, use_bias=True, name=f"{name}_conv_reduce"
+    )(x)
+    if activation is not None:
+        x = layers.Activation(activation, name=f"{name}_act1")(x)
+    x = layers.Conv2D(
+        input_channels, 1, use_bias=True, name=f"{name}_conv_expand"
+    )(x)
+    if activation is not None:
+        x = layers.Activation(gate_activation, name=f"{name}_gate")(x)
+    x = layers.Multiply(name=name)([inputs, x])
+    return x

kimm/blocks/transformer_block.py → kimm/_src/blocks/transformer.py

@@ -3,9 +3,11 @@
 from keras import backend
 from keras import layers
 
-from kimm import layers as kimm_layers
+from kimm._src.kimm_export import kimm_export
+from kimm._src.layers.attention import Attention
 
 
+@kimm_export(parent_path=["kimm.blocks"])
 def apply_mlp_block(
     inputs,
     hidden_dim: int,
@@ -42,6 +44,7 @@ def apply_mlp_block(
     return x
 
 
+@kimm_export(parent_path=["kimm.blocks"])
 def apply_transformer_block(
     inputs,
     dim: int,
@@ -58,7 +61,7 @@ def apply_transformer_block(
     residual_1 = x
 
     x = layers.LayerNormalization(epsilon=1e-6, name=f"{name}_norm1")(x)
-    x = kimm_layers.Attention(
+    x = Attention(
         dim,
         num_heads,
         use_qkv_bias,

kimm/export/export_onnx.py → kimm/_src/export/export_onnx.py

@@ -6,10 +6,12 @@
 from keras import models
 from keras import ops
 
-from kimm.models import BaseModel
-from kimm.utils.module_utils import torch
+from kimm._src.kimm_export import kimm_export
+from kimm._src.models.base_model import BaseModel
+from kimm._src.utils.module_utils import torch
 
 
+@kimm_export(parent_path=["kimm.export"])
 def export_onnx(
     model: BaseModel,
     input_shape: typing.Union[int, typing.Sequence[int]],

kimm/export/export_onnx_test.py → kimm/_src/export/export_onnx_test.py

@@ -3,14 +3,16 @@
 from keras import backend
 from keras.src import testing
 
-from kimm import export
-from kimm import models
+from kimm._src import models
+from kimm._src.export import export_onnx
 
 
 class ExportOnnxTest(testing.TestCase, parameterized.TestCase):
     def get_model(self):
         input_shape = [3, 224, 224]  # channels_first
-        model = models.MobileNetV3W050Small(include_preprocessing=False)
+        model = models.mobilenet_v3.MobileNetV3W050Small(
+            include_preprocessing=False, weights=None
+        )
         return input_shape, model
 
     @classmethod
@@ -33,4 +35,4 @@ def DISABLE_test_export_onnx_use(self):
 
         temp_dir = self.get_temp_dir()
 
-        export.export_onnx(model, input_shape, f"{temp_dir}/model.onnx")
+        export_onnx.export_onnx(model, input_shape, f"{temp_dir}/model.onnx")

kimm/export/export_tflite.py → kimm/_src/export/export_tflite.py

@@ -7,9 +7,11 @@
 from keras import models
 from keras.src.utils.module_utils import tensorflow as tf
 
-from kimm.models import BaseModel
+from kimm._src.kimm_export import kimm_export
+from kimm._src.models.base_model import BaseModel
 
 
+@kimm_export(parent_path=["kimm.export"])
 def export_tflite(
     model: BaseModel,
     input_shape: typing.Union[int, typing.Sequence[int]],
@@ -20,9 +22,10 @@ def export_tflite(
 ):
     """Export the model to tflite format.
 
-    Only tensorflow backend with 'channels_last' is supported. The tflite model
-    will be generated using `tf.lite.TFLiteConverter.from_saved_model` and
-    optimized through tflite built-in functions.
+    Only TensorFlow backend with 'channels_last' is supported. The
+    tflite model will be generated using
+    `tf.lite.TFLiteConverter.from_saved_model` and optimized through tflite
+    built-in functions.
 
     Note that when exporting an `int8` tflite model, `representative_dataset`
     must be passed.
@@ -37,8 +40,8 @@ def export_tflite(
         batch_size: int, specifying the batch size of the input,
             defaults to `1`.
     """
-    if backend.backend() != "tensorflow":
-        raise ValueError("`export_tflite` only supports tensorflow backend")
+    if backend.backend() not in ("tensorflow",):
+        raise ValueError("`export_tflite` only supports TensorFlow backend")
     if backend.image_data_format() != "channels_last":
         raise ValueError(
             "`export_tflite` only supports 'channels_last' data format."