feat: Implement the conv2D rewrites for int8 and fp32 models

Enable clustering and fully connected rewrites for conv2D layers.

Resolves: MLIA-1159 and MLIA-1160

Signed-off-by: Nathan Bailey <nathan.bailey@arm.com>
Change-Id: I640b8a7e79e455b12fb68d02ac1c33213b8de9c6

4 files changed, 67 insertions, 16 deletions

diff --git a/src/mlia/nn/rewrite/core/rewrite.py b/src/mlia/nn/rewrite/core/rewrite.py
index 8fba806..6674d02 100644
--- a/src/mlia/nn/rewrite/core/rewrite.py
+++ b/src/mlia/nn/rewrite/core/rewrite.py
@@ -15,6 +15,9 @@ from typing import Callable
 import numpy as np
 import tensorflow_model_optimization as tfmot
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_utils import (  # pylint: disable=no-name-in-module
+    is_pruned_m_by_n,
+)
 
 from mlia.core.errors import ConfigurationError
 from mlia.core.reporting import Column
@@ -32,7 +35,6 @@ from mlia.nn.rewrite.library.sparsity import fc_sparsity_rewrite
 from mlia.nn.tensorflow.config import TFLiteModel
 from mlia.utils.registry import Registry
 
-
 logger = logging.getLogger(__name__)
 RewriteCallable = Callable[[Any, Any], keras.Model]
 
@@ -131,7 +133,20 @@ class Sparsity24Rewrite(QuantizeAwareTrainingRewrite):
         return model
 
     def check_optimization(self, model: keras.Model, **kwargs: Any) -> bool:
-        """Not needed here."""
+        """Check if sparity has produced the correct result."""
+        for layer in model.layers:
+            for weight in layer.weights:
+                if "kernel" in weight.name:
+                    if "kernel_min" in weight.name or "kernel_max" in weight.name:
+                        continue
+                    if not is_pruned_m_by_n(weight, m_by_n=(2, 4)):
+                        logger.warning(
+                            "\nWARNING: Could not find (2,4) sparsity, "
+                            "in layer %s for weight %s \n",
+                            layer.name,
+                            weight.name,
+                        )
+                        return False
         return True
 
 
@@ -251,12 +266,6 @@ class RewritingOptimizer(Optimizer):
             self.optimizer_configuration.optimization_target
         ]
 
-        if self.optimizer_configuration.optimization_target in [
-            "conv2d-clustering",
-            "conv2d-sparsity24",
-        ]:
-            raise NotImplementedError
-
         use_unmodified_model = True
         tflite_model = self.model.model_path
         tfrecord = str(self.optimizer_configuration.dataset)
diff --git a/src/mlia/nn/rewrite/library/clustering.py b/src/mlia/nn/rewrite/library/clustering.py
index 1d22dc6..6f06c48 100644
--- a/src/mlia/nn/rewrite/library/clustering.py
+++ b/src/mlia/nn/rewrite/library/clustering.py
@@ -6,11 +6,13 @@ from typing import Any
 import tensorflow_model_optimization as tfmot
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
+from mlia.nn.rewrite.library.helper_functions import compute_conv2d_parameters
+
 
 def fc_clustering_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for clustering rewrite."""
     rewrite_params = {
-        "number_of_clusters": 32,
+        "number_of_clusters": 4,
         "cluster_centroids_init": tfmot.clustering.keras.CentroidInitialization.LINEAR,
     }
     model = tfmot.clustering.keras.cluster_weights(
@@ -26,19 +28,22 @@ def fc_clustering_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     return model
 
 
-def conv2d_clustering_rewrite(
-    input_shape: Any, conv2d_parameters: dict[str, Any]
-) -> keras.Model:
+def conv2d_clustering_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for clustering rewrite."""
     rewrite_params = {
-        "number_of_clusters": 32,
+        "number_of_clusters": 4,
         "cluster_centroids_init": tfmot.clustering.keras.CentroidInitialization.LINEAR,
     }
+    conv2d_parameters = compute_conv2d_parameters(
+        input_shape=input_shape, output_shape=output_shape
+    )
     model = tfmot.clustering.keras.cluster_weights(
         to_cluster=keras.Sequential(
             [
                 keras.layers.InputLayer(input_shape=input_shape),
                 keras.layers.Conv2D(**conv2d_parameters),
+                keras.layers.BatchNormalization(),
+                keras.layers.ReLU(),
             ]
         ),
         **rewrite_params
diff --git a/src/mlia/nn/rewrite/library/helper_functions.py b/src/mlia/nn/rewrite/library/helper_functions.py
new file mode 100644
index 0000000..4f08170
--- /dev/null
+++ b/src/mlia/nn/rewrite/library/helper_functions.py
@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""Helper functions for the rewrite library."""
+import math
+from typing import Any
+
+import numpy as np
+
+
+def compute_conv2d_parameters(
+    input_shape: np.ndarray, output_shape: np.ndarray
+) -> dict[str, Any]:
+    """Compute needed kernel size and strides for a given input and output_shape."""
+    input_shape = input_shape.tolist()
+    output_shape = output_shape.tolist()
+    assert len(input_shape) == 3
+    assert len(output_shape) == 3
+    num_filters = (output_shape[-1] - input_shape[-1]) + input_shape[-1]
+    padding = "valid"
+    kernel_size = (3, 3)
+    stride_h = round(input_shape[0] / output_shape[0])
+    check_output_size_h = math.floor((input_shape[0] - kernel_size[0]) / stride_h) + 1
+    stride_w = round(input_shape[1] / output_shape[1])
+    check_output_size_w = math.floor((input_shape[1] - kernel_size[1]) / stride_w) + 1
+    if check_output_size_h != output_shape[0] or check_output_size_w != output_shape[1]:
+        padding = "same"
+    return {
+        "filters": num_filters,
+        "kernel_size": kernel_size,
+        "padding": padding,
+        "strides": (stride_h, stride_w),
+    }
diff --git a/src/mlia/nn/rewrite/library/sparsity.py b/src/mlia/nn/rewrite/library/sparsity.py
index 8b74b72..709593a 100644
--- a/src/mlia/nn/rewrite/library/sparsity.py
+++ b/src/mlia/nn/rewrite/library/sparsity.py
@@ -6,6 +6,8 @@ from typing import Any
 import tensorflow_model_optimization as tfmot
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
+from mlia.nn.rewrite.library.helper_functions import compute_conv2d_parameters
+
 
 def fc_sparsity_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for rewrite."""
@@ -23,15 +25,18 @@ def fc_sparsity_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     return model
 
 
-def conv2d_sparsity_rewrite(
-    input_shape: Any, conv2d_parameters: dict[str, Any]
-) -> keras.Model:
+def conv2d_sparsity_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for rewrite."""
+    conv2d_parameters = compute_conv2d_parameters(
+        input_shape=input_shape, output_shape=output_shape
+    )
     model = tfmot.sparsity.keras.prune_low_magnitude(
         to_prune=keras.Sequential(
             [
                 keras.layers.InputLayer(input_shape=input_shape),
                 keras.layers.Conv2D(**conv2d_parameters),
+                keras.layers.BatchNormalization(),
+                keras.layers.ReLU(),
             ]
         ),
         sparsity_m_by_n=(2, 4),