feat: CLI and API changes for the conv2d rewrites

Implements CLI and API changes for the new conv2d rewrite targets

Resolves: MLIA-1157

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

4 files changed, 55 insertions, 12 deletions

diff --git a/src/mlia/nn/rewrite/core/rewrite.py b/src/mlia/nn/rewrite/core/rewrite.py
index e2c097c..8fba806 100644
--- a/src/mlia/nn/rewrite/core/rewrite.py
+++ b/src/mlia/nn/rewrite/core/rewrite.py
@@ -24,13 +24,11 @@ from mlia.nn.common import Optimizer
 from mlia.nn.common import OptimizerConfiguration
 from mlia.nn.rewrite.core.train import train
 from mlia.nn.rewrite.core.train import TrainingParameters
-from mlia.nn.rewrite.library.fc_clustering_layer import (
-    get_keras_model_clus as fc_clustering_rewrite,
-)
-from mlia.nn.rewrite.library.fc_layer import get_keras_model as fc_rewrite
-from mlia.nn.rewrite.library.fc_sparsity24_layer import (
-    get_keras_model as fc_rewrite_sparsity24,
-)
+from mlia.nn.rewrite.library.clustering import conv2d_clustering_rewrite
+from mlia.nn.rewrite.library.clustering import fc_clustering_rewrite
+from mlia.nn.rewrite.library.fc_layer import fc_rewrite
+from mlia.nn.rewrite.library.sparsity import conv2d_sparsity_rewrite
+from mlia.nn.rewrite.library.sparsity import fc_sparsity_rewrite
 from mlia.nn.tensorflow.config import TFLiteModel
 from mlia.utils.registry import Registry
 
@@ -227,8 +225,10 @@ class RewritingOptimizer(Optimizer):
     registry = RewriteRegistry(
         [
             GenericRewrite("fully-connected", fc_rewrite),
-            Sparsity24Rewrite("fully-connected-sparsity24", fc_rewrite_sparsity24),
+            Sparsity24Rewrite("fully-connected-sparsity24", fc_sparsity_rewrite),
             ClusteringRewrite("fully-connected-clustering", fc_clustering_rewrite),
+            ClusteringRewrite("conv2d-clustering", conv2d_clustering_rewrite),
+            Sparsity24Rewrite("conv2d-sparsity24", conv2d_sparsity_rewrite),
         ]
     )
 
@@ -251,6 +251,12 @@ 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/fc_clustering_layer.py b/src/mlia/nn/rewrite/library/clustering.py
index 7cc383e..1d22dc6 100644
--- a/src/mlia/nn/rewrite/library/fc_clustering_layer.py
+++ b/src/mlia/nn/rewrite/library/clustering.py
@@ -7,7 +7,7 @@ import tensorflow_model_optimization as tfmot
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
 
-def get_keras_model_clus(input_shape: Any, output_shape: Any) -> keras.Model:
+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,
@@ -24,3 +24,23 @@ def get_keras_model_clus(input_shape: Any, output_shape: Any) -> keras.Model:
         **rewrite_params
     )
     return model
+
+
+def conv2d_clustering_rewrite(
+    input_shape: Any, conv2d_parameters: dict[str, Any]
+) -> keras.Model:
+    """Generate TensorFlow Lite model for clustering rewrite."""
+    rewrite_params = {
+        "number_of_clusters": 32,
+        "cluster_centroids_init": tfmot.clustering.keras.CentroidInitialization.LINEAR,
+    }
+    model = tfmot.clustering.keras.cluster_weights(
+        to_cluster=keras.Sequential(
+            [
+                keras.layers.InputLayer(input_shape=input_shape),
+                keras.layers.Conv2D(**conv2d_parameters),
+            ]
+        ),
+        **rewrite_params
+    )
+    return model
diff --git a/src/mlia/nn/rewrite/library/fc_layer.py b/src/mlia/nn/rewrite/library/fc_layer.py
index 041ce85..cb98cb9 100644
--- a/src/mlia/nn/rewrite/library/fc_layer.py
+++ b/src/mlia/nn/rewrite/library/fc_layer.py
@@ -6,7 +6,7 @@ from typing import Any
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
 
-def get_keras_model(input_shape: Any, output_shape: Any) -> keras.Model:
+def fc_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for rewrite."""
     model = keras.Sequential(
         (
diff --git a/src/mlia/nn/rewrite/library/fc_sparsity24_layer.py b/src/mlia/nn/rewrite/library/sparsity.py
index 531b34a..8b74b72 100644
--- a/src/mlia/nn/rewrite/library/fc_sparsity24_layer.py
+++ b/src/mlia/nn/rewrite/library/sparsity.py
@@ -1,13 +1,13 @@
 # SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates.
 # SPDX-License-Identifier: Apache-2.0
-"""Example rewrite with one fully connected 2:4 sparsity layer."""
+"""Example rewrite with one fully connected clustered layer."""
 from typing import Any
 
 import tensorflow_model_optimization as tfmot
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
 
-def get_keras_model(input_shape: Any, output_shape: Any) -> keras.Model:
+def fc_sparsity_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for rewrite."""
     model = tfmot.sparsity.keras.prune_low_magnitude(
         to_prune=keras.Sequential(
@@ -21,3 +21,20 @@ def get_keras_model(input_shape: Any, output_shape: Any) -> keras.Model:
     )
 
     return model
+
+
+def conv2d_sparsity_rewrite(
+    input_shape: Any, conv2d_parameters: dict[str, Any]
+) -> keras.Model:
+    """Generate TensorFlow Lite model for rewrite."""
+    model = tfmot.sparsity.keras.prune_low_magnitude(
+        to_prune=keras.Sequential(
+            [
+                keras.layers.InputLayer(input_shape=input_shape),
+                keras.layers.Conv2D(**conv2d_parameters),
+            ]
+        ),
+        sparsity_m_by_n=(2, 4),
+    )
+
+    return model