Add LCNet and fix model serialization (#10)

* Merge `apply_depthwise_separation_block`

* Add `LCNet`

* Speed up gpu test

* Cleanup

* Update `add_model_to_registry`

* Fix model serialization

conftest.py

@@ -2,6 +2,13 @@
 
 
 def pytest_configure():
+    import tensorflow as tf
+
+    # disable tensorflow gpu memory preallocation
+    physical_devices = tf.config.list_physical_devices("GPU")
+    for device in physical_devices:
+        tf.config.experimental.set_memory_growth(device, True)
+
     # disable jax gpu memory preallocation
     # https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
     os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

kimm/blocks/__init__.py

@@ -1,6 +1,9 @@
 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
+from kimm.blocks.depthwise_separation_block import (
+    apply_depthwise_separation_block,
+)
 from kimm.blocks.inverted_residual_block import apply_inverted_residual_block
 from kimm.blocks.transformer_block import apply_mlp_block
 from kimm.blocks.transformer_block import apply_transformer_block

kimm/blocks/depthwise_separation_block.py

@@ -0,0 +1,59 @@
+from keras import layers
+
+from kimm.blocks.base_block import apply_conv2d_block
+from kimm.blocks.base_block import apply_se_block
+
+
+def apply_depthwise_separation_block(
+    inputs,
+    output_channels,
+    depthwise_kernel_size=3,
+    pointwise_kernel_size=1,
+    strides=1,
+    se_ratio=0.0,
+    activation="swish",
+    se_activation="relu",
+    se_gate_activation="sigmoid",
+    se_make_divisible_number=None,
+    pw_activation=None,
+    skip=True,
+    bn_epsilon=1e-5,
+    padding=None,
+    name="depthwise_separation_block",
+):
+    input_channels = inputs.shape[-1]
+    has_skip = skip and (strides == 1 and input_channels == output_channels)
+
+    x = inputs
+    x = apply_conv2d_block(
+        x,
+        kernel_size=depthwise_kernel_size,
+        strides=strides,
+        activation=activation,
+        use_depthwise=True,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_dw",
+    )
+    if se_ratio > 0:
+        x = apply_se_block(
+            x,
+            se_ratio,
+            activation=se_activation,
+            gate_activation=se_gate_activation,
+            make_divisible_number=se_make_divisible_number,
+            name=f"{name}_se",
+        )
+    x = apply_conv2d_block(
+        x,
+        output_channels,
+        pointwise_kernel_size,
+        1,
+        activation=pw_activation,
+        bn_epsilon=bn_epsilon,
+        padding=padding,
+        name=f"{name}_conv_pw",
+    )
+    if has_skip:
+        x = layers.Add()([x, inputs])
+    return x