Init

Author: james77777778

.github/dependabot.yml

@@ -0,0 +1,23 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "monthly"
+    groups:
+      github-actions:
+        patterns:
+          - "*"
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "monthly"
+    groups:
+      python:
+        patterns:
+          - "*"

.github/workflows/actions.yml

@@ -0,0 +1,73 @@
+# Ref: https://github.com/keras-team/keras/blob/master/.github/workflows/actions.yml
+name: Tests
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+  release:
+    types: [created]
+
+permissions:
+  contents: read
+
+jobs:
+  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 }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - 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@v3
+        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: Test with pytest
+        run: |
+          pytest
+
+  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@v3
+        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

.gitignore

@@ -158,3 +158,7 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Keras
+*.keras
+exported

README.md

@@ -1 +1,10 @@
-# kimm
+# Keras Image Models
+
+## Unit Tests
+
+```bash
+# KERAS_BACKEND=jax|numpy|tensorflow|torch
+CUDA_VISIBLE_DEVICES= KERAS_BACKEND=tensorflow pytest
+```
+
+## Acknowledgments

conftest.py

@@ -0,0 +1,7 @@
+import os
+
+
+def pytest_configure():
+    # disable jax gpu memory preallocation
+    # https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
+    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

kimm/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

kimm/blocks/__init__.py

@@ -0,0 +1,3 @@
+from kimm.blocks.base_block import apply_activation
+from kimm.blocks.base_block import apply_conv2d_block
+from kimm.blocks.base_block import apply_se_block

kimm/blocks/base_block.py

@@ -0,0 +1,91 @@
+from keras import layers
+
+from kimm.utils import make_divisible
+
+
+def apply_activation(x, activation=None, name="activation"):
+    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
+
+
+def apply_conv2d_block(
+    inputs,
+    filters,
+    kernel_size,
+    strides=1,
+    groups=1,
+    activation=None,
+    use_depthwise=False,
+    bn_momentum=0.9,
+    bn_epsilon=1e-5,
+    name="conv2d_block",
+):
+    x = inputs
+
+    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(
+            filters,
+            kernel_size,
+            strides,
+            padding=padding,
+            groups=groups,
+            use_bias=False,
+            name=f"{name}_conv2d",
+        )(x)
+    else:
+        x = layers.DepthwiseConv2D(
+            kernel_size,
+            strides,
+            padding=padding,
+            use_bias=False,
+            name=f"{name}_dwconv2d",
+        )(x)
+    x = layers.BatchNormalization(
+        name=f"{name}_bn", momentum=bn_momentum, epsilon=bn_epsilon
+    )(x)
+    x = apply_activation(x, activation, name=name)
+    return x
+
+
+def apply_se_block(
+    inputs,
+    se_ratio=0.25,
+    activation="relu",
+    gate_activation="sigmoid",
+    make_divisible_number=None,
+    name="se_block",
+):
+    input_channels = inputs.shape[-1]
+    if make_divisible_number is None:
+        se_channels = round(input_channels * se_ratio)
+    else:
+        se_channels = make_divisible(
+            input_channels * se_ratio, make_divisible_number
+        )
+
+    ori_x = inputs
+    x = inputs
+    x = layers.GlobalAveragePooling2D(keepdims=True, name=f"{name}_mean")(x)
+    x = layers.Conv2D(
+        se_channels, 1, use_bias=True, name=f"{name}_reduce_conv2d"
+    )(x)
+    x = apply_activation(x, activation, name=f"{name}_act")
+    x = layers.Conv2D(
+        input_channels, 1, use_bias=True, name=f"{name}_expand_conv2d"
+    )(x)
+    x = apply_activation(x, gate_activation, name=f"{name}_gate_act")
+    out = layers.Multiply(name=name)([ori_x, x])
+    return out

kimm/layers/__init__.py

@@ -0,0 +1,3 @@
+from kimm.layers.attention import Attention
+from kimm.layers.layer_scale import LayerScale
+from kimm.layers.position_embedding import PositionEmbedding

kimm/layers/attention.py

@@ -0,0 +1,105 @@
+from keras import layers
+from keras import ops
+
+
+class Attention(layers.Layer):
+    def __init__(
+        self,
+        hidden_dim,
+        num_heads=8,
+        use_qkv_bias=False,
+        use_qk_norm=False,
+        attention_dropout_rate=0.0,
+        projection_dropout_rate=0.0,
+        name="attention",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+        self.scale = self.head_dim ** (-0.5)
+        self.use_qkv_bias = use_qkv_bias
+        self.use_qk_norm = use_qk_norm
+        self.attention_dropout_rate = attention_dropout_rate
+        self.projection_dropout_rate = projection_dropout_rate
+        self.name = name
+
+        self.qkv = layers.Dense(
+            hidden_dim * 3, use_bias=use_qkv_bias, 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")
+        else:
+            self.q_norm = layers.Identity()
+            self.k_norm = layers.Identity()
+
+        self.attention_dropout = layers.Dropout(
+            attention_dropout_rate, name=f"{name}_attn_drop"
+        )
+        self.projection = layers.Dense(hidden_dim, name=f"{name}_proj")
+        self.projection_dropout = layers.Dropout(
+            projection_dropout_rate, name=f"{name}_proj_drop"
+        )
+
+    def call(self, inputs, training=None, mask=None):
+        input_shape = ops.shape(inputs)
+        qkv = self.qkv(inputs)
+        qkv = ops.reshape(
+            qkv,
+            [
+                input_shape[0],
+                input_shape[1],
+                3,
+                self.num_heads,
+                self.head_dim,
+            ],
+        )
+        qkv = ops.transpose(qkv, [2, 0, 3, 1, 4])
+        q, k, v = ops.unstack(qkv, 3, axis=0)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # attention
+        q = ops.multiply(q, self.scale)
+        attn = ops.matmul(q, ops.swapaxes(k, -2, -1))
+        attn = ops.softmax(attn)
+        attn = self.attention_dropout(attn)
+        x = ops.matmul(attn, v)
+
+        x = ops.swapaxes(x, 1, 2)
+        x = ops.reshape(x, input_shape)
+        x = self.projection(x)
+        x = self.projection_dropout(x)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "num_heads": self.num_heads,
+                "use_qkv_bias": self.use_qkv_bias,
+                "use_qk_norm": self.use_qk_norm,
+                "attention_dropout_rate": self.attention_dropout_rate,
+                "projection_dropout_rate": self.projection_dropout_rate,
+                "name": self.name,
+            }
+        )
+        return config
+
+
+if __name__ == "__main__":
+    from keras import models
+    from keras import random
+
+    inputs = layers.Input(shape=[197, 768])
+    outputs = Attention(768)(inputs)
+
+    model = models.Model(inputs, outputs)
+    model.summary()
+
+    inputs = random.uniform([1, 197, 768])
+    outputs = model(inputs)
+    print(outputs.shape)