optional dependencies, fix depthwise bug

Author: nroope

pyproject.toml

@@ -23,7 +23,11 @@ classifiers = [
   "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dynamic = [ "version" ]
-dependencies = [ "keras>=3", "pyyaml>=6.0.1", "quantizers>=1.1", "torch>=2.1", "pydantic>=2.0"]
+dependencies = [ "hgq2", "keras>=3", "optuna", "pydantic>=2", "pyyaml>=6.0.1", "quantizers>=1.1" ]
+optional-dependencies.all = [ "pytest>=8.4", "tensorflow>=2.17,<=2.20", "torch>=2.1" ]
+optional-dependencies.tensorflow = [ "tensorflow>=2.17,<=2.20" ]
+optional-dependencies.test = [ "pytest>=8.4" ]
+optional-dependencies.torch = [ "torch>=2.1" ]
 urls.repository = "https://github.com/nroope/PQuant"
 
 [tool.setuptools]

src/pquant/core/keras/layers.py

@@ -16,6 +16,7 @@
     ReLU,
     SeparableConv2D,
 )
+from keras.src.layers.input_spec import InputSpec
 from keras.src.ops.operation_utils import compute_pooling_output_shape
 
 from pquant.core.keras.activations import PQActivation
@@ -276,10 +277,8 @@ def __init__(
         self.use_bias = use_bias
         self.strides = strides
         self.dilation_rate = dilation_rate
-        # self.weight_transpose = (2, 3, 0, 1)
-        # self.weight_transpose_back = (2, 3, 1, 0)
-        self.weight_transpose = (3, 2, 0, 1)
-        self.weight_transpose_back = (2, 3, 1, 0)
+        self.weight_transpose = (2, 3, 0, 1)
+        self.weight_transpose_back = (2, 3, 0, 1)
         self.data_transpose = (0, 3, 1, 2)
         self.do_transpose_data = self.data_format == "channels_last"
         self._weight = None
@@ -288,10 +287,12 @@ def __init__(
     def build(self, input_shape):
         super().build(input_shape)
         if self.data_format == "channels_last":
+            channel_axis = -1
             input_channel = input_shape[-1]
         else:
+            channel_axis = 1
             input_channel = input_shape[1]
-
+        self.input_spec = InputSpec(min_ndim=self.rank + 2, axes={channel_axis: input_channel})
         depthwise_shape = self.kernel_size + (
             input_channel,
             self.depth_multiplier,

tests/test_keras_compression_layers.py

@@ -735,15 +735,18 @@ def test_ap_conv1d_channels_last_transpose(config_ap, conv1d_input):
 
 
 def test_ap_depthwiseconv2d_channels_last_transpose(config_ap, conv2d_input):
+    if keras.backend.image_data_format() == "channels_last":
+        conv2d_input = ops.transpose(conv2d_input, (0, 3, 1, 2))
     keras.backend.set_image_data_format("channels_first")
     inp = ops.reshape(ops.linspace(0, 1, ops.size(conv2d_input)), conv2d_input.shape)
 
     inputs = keras.Input(shape=inp.shape[1:])
-    out = DepthwiseConv2D(KERNEL_SIZE, use_bias=False, padding="same")(inputs)
+    out = DepthwiseConv2D(KERNEL_SIZE, use_bias=False, padding="same", data_format="channels_first")(inputs)
     model_cf = keras.Model(inputs=inputs, outputs=out, name="test_dwconv2d")
     model_cf = add_compression_layers(model_cf, config_ap, inp.shape)
-    weight_cf = model_cf.layers[1]._kernel
 
+    weight_cf = model_cf.layers[1]._kernel
+    model_cf.summary()
     post_pretrain_functions(model_cf, config_ap)
     model_cf(inp, training=True)
     model_cf(inp, training=True)