diff options
author | Nathan Bailey <nathan.bailey@arm.com> | 2024-05-15 08:12:30 +0100 |
---|---|---|
committer | Nathan Bailey <nathan.bailey@arm.com> | 2024-05-21 16:51:15 +0100 |
commit | 856111bcaef76c60303bdf2ae7cbf718d93d1df4 (patch) | |
tree | d955901817194e48e478f751140bd3c1741d1834 /src/mlia/nn | |
parent | 0d3cc76284f9311c99169b568570d767f5b0aeb6 (diff) | |
download | mlia-856111bcaef76c60303bdf2ae7cbf718d93d1df4.tar.gz |
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
Diffstat (limited to 'src/mlia/nn')
-rw-r--r-- | src/mlia/nn/rewrite/core/rewrite.py | 25 | ||||
-rw-r--r-- | src/mlia/nn/rewrite/library/clustering.py | 15 | ||||
-rw-r--r-- | src/mlia/nn/rewrite/library/helper_functions.py | 32 | ||||
-rw-r--r-- | src/mlia/nn/rewrite/library/sparsity.py | 11 |
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), |