Repository: quic/aimet (language: python)

Issue:
Add SVD Splitter for keras



Discussion / hints:

Fix patch:
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/svd_pruner.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/svd_pruner.py
new file mode 100644
index 0000000000..e760faca85
--- /dev/null
+++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/svd_pruner.py
@@ -0,0 +1,78 @@
+# /usr/bin/env python3.8
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2022, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+
+""" Prunes layers using Spatial SVD schemes """
+
+from aimet_common.utils import AimetLogger
+import aimet_common.svd_pruner
+
+from aimet_tensorflow.keras.layer_database import LayerDatabase, Layer
+from aimet_tensorflow.keras.svd_spiltter import SpatialSvdModuleSplitter
+
+logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Svd)
+
+
+class SpatialSvdPruner(aimet_common.svd_pruner.SpatialSvdPruner):
+    """
+    Pruner for Spatial-SVD method
+    """
+
+    def _perform_svd_and_split_layer(self, layer: Layer, rank: int, comp_layer_db: LayerDatabase):
+
+        """
+        Performs spatial svd and splits given layer into two layers
+        :param layer: Layer to split
+        :param rank: Rank to use for spatial svd splitting
+        :param comp_layer_db: Compressed layer db to update with the split layers
+        :return: None
+        """
+
+        # Split module using Spatial SVD
+        module_a, module_b = SpatialSvdModuleSplitter.split_module(comp_layer_db.model, layer, rank)
+
+        # get the output activation shape for first conv op
+        output_shape_a = module_a.output_shape
+
+        # get the output activation shape for second conv op
+        output_shape_b = module_b.output_shape
+
+        # Create two new layers and return them
+        layer_a = Layer(layer=module_a, name=module_a.name, output_shape=output_shape_a)
+        layer_b = Layer(layer=module_b, name=module_b.name, output_shape=output_shape_b)
+
+        comp_layer_db.replace_layer_with_sequential_of_two_layers(layer, layer_a, layer_b)
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/svd_spiltter.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/svd_spiltter.py
new file mode 100644
index 0000000000..8ff2ef5c15
--- /dev/null
+++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/svd_spiltter.py
@@ -0,0 +1,132 @@
+# /usr/bin/env python3.8
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2022, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+
+""" Implementation of layer splitting logic for spatial svd schemes """
+import numpy as np
+import tensorflow as tf
+
+from aimet_tensorflow.keras.utils.op.conv import get_strides_for_split_conv_ops
+from aimet_tensorflow.keras.layer_database import Layer
+from aimet_tensorflow.keras.utils.common import replace_layer_in_functional_model
+from aimet_tensorflow.keras.utils.weight_tensor_utils import WeightTensorUtils
+from aimet_common.utils import AimetLogger
+from aimet_common.svd_pruner import SpatialSvdPruner
+
+logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Svd)
+
+
+class SpatialSvdModuleSplitter:
+    """ Spatial SVD module splitter"""
+
+    @staticmethod
+    def split_module(model: tf.keras.Model, layer: Layer, rank: int) -> (tf.keras.layers.Layer, tf.keras.layers.Layer):
+        """
+        :param  model: Keras Model whose layer we want to split
+        :param layer: Module to be split
+        :param rank: rank for splitting
+        :return: Two split modules
+        """
+
+        h, v = SpatialSvdModuleSplitter.get_svd_matrices(layer, rank)
+
+        conv_a_stride, conv_b_stride = get_strides_for_split_conv_ops(layer=layer.module)
+
+        data_format_channels = layer.module.data_format
+        padding = layer.module.padding
+
+        conv_a = tf.keras.layers.Conv2D(filters=v.shape[3], kernel_size=(v.shape[0], v.shape[1]),
+                                        strides=conv_a_stride, data_format=data_format_channels,
+                                        activation=None, padding=padding,
+                                        name=layer.module.name + '_a', use_bias=False)
+
+        # get the succeeding bias tensor if present
+        use_bias = False
+        if len(layer.module.get_weights()) > 1:
+            use_bias = True
+
+        conv_b = tf.keras.layers.Conv2D(filters=h.shape[3], kernel_size=(h.shape[0], h.shape[1]),
+                                        strides=conv_b_stride,
+                                        name=layer.module.name + '_b',
+                                        data_format=data_format_channels, padding=padding, use_bias=use_bias)
+
+        # Replace the layer in the model
+        replace_layer_in_functional_model(model, layer.module, [conv_a, conv_b])
+
+        # Check if the weight shape are equal or not
+        assert conv_a.get_weights()[0].shape == v.shape
+        assert conv_b.get_weights()[0].shape == h.shape
+
+        # Set the weights (kernel) for conv_a
+        conv_a.set_weights([v])
+
+        # Set the weights (kernel and bias )for conv_b
+        conv_b_weight_tensor = [h]
+
+        if use_bias:
+            bias_tensor = layer.module.get_weights()[1]
+            conv_b_weight_tensor.append(bias_tensor)
+
+        conv_b.set_weights(conv_b_weight_tensor)
+
+        return conv_a, conv_b
+
+    @staticmethod
+    def get_svd_matrices(layer: Layer, rank: int) -> (np.array, np.array):
+        """
+        :param layer: Module to be split
+        :param rank: rank for splitting
+        :return: v and h matrices after Single Value Decomposition
+        """
+
+        # get the weight parameters
+        weight_tensor = layer.module.get_weights()[0]
+
+        # Conv2d weight shape in TensorFlow  [kh, kw, Nic, Noc]
+        # re order in the common shape  [Noc, Nic, kh, kw]
+        weight_tensor = WeightTensorUtils.transpose_from_tf_to_libpymo_format(weight_tensor, layer.module)
+
+        out_channels, in_channels, height, width = weight_tensor.shape
+
+        h, v = SpatialSvdPruner.lingalg_spatial_svd(weight_tensor, rank, in_channels, out_channels, height, width)
+
+        # h, v matrices are in the common shape [Noc, Nic, kh, kw]
+        # re order in TensorFlow Conv2d shape [kh, kw, Nic, Noc]
+        h = WeightTensorUtils.transpose_from_libpymo_to_tf_format(h, layer.module)
+        v = WeightTensorUtils.transpose_from_libpymo_to_tf_format(v, layer.module)
+
+        return h, v


Test patch:
diff --git a/TrainingExtensions/tensorflow/test/python/eager/test_spatial_svd.py b/TrainingExtensions/tensorflow/test/python/eager/test_spatial_svd.py
new file mode 100644
index 0000000000..c79738a493
--- /dev/null
+++ b/TrainingExtensions/tensorflow/test/python/eager/test_spatial_svd.py
@@ -0,0 +1,158 @@
+# /usr/bin/env python3.8
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2022, Qualcomm Innovation Center, Inc. All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  1. Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#  2. Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#  3. Neither the name of the copyright holder nor the names of its contributors
+#     may be used to endorse or promote products derived from this software
+#     without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+#  POSSIBILITY OF SUCH DAMAGE.
+#
+#  SPDX-License-Identifier: BSD-3-Clause
+#
+#  @@-COPYRIGHT-END-@@
+# =============================================================================
+
+# import unittest
+import pytest
+
+import numpy as np
+from aimet_common.utils import AimetLogger
+from aimet_tensorflow.keras.layer_database import *
+from aimet_tensorflow.keras.svd_spiltter import SpatialSvdModuleSplitter
+from aimet_tensorflow.keras.svd_pruner import SpatialSvdPruner
+
+logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Test)
+
+def get_model(model_type = "Sequential"):
+    tf.keras.backend.clear_session()
+    if model_type == "Sequential":
+        return  tf.keras.Sequential([
+            tf.keras.layers.Reshape(target_shape=(28, 28, 1), input_shape=(28 * 28,)),
+            tf.keras.layers.Conv2D(32, 5, strides=(2, 2), name='conv1', padding='same'),
+            tf.keras.layers.Conv2D(64, 32, name='conv2', padding='same'),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(10, name='linear')
+        ])
+    elif model_type == "Functional":
+        inp = tf.keras.Input((28*28))
+        x = tf.keras.layers.Reshape(target_shape=(28, 28, 1))(inp)
+        x = tf.keras.layers.Conv2D(32, 5, strides=(2, 2), name='conv1', padding='same')(x)
+        x = tf.keras.layers.Conv2D(64, 32, name='conv2', padding='same')(x)
+        x = tf.keras.layers.Flatten()(x)
+        out = tf.keras.layers.Dense(10, name = 'linear')(x)
+
+        return tf.keras.Model(inp, out)
+
+def _get_layers(model, model_type="Sequential"):
+    # Drop first layer (Input layer) of Functional model
+    if model_type == "Functional":
+        return model.layers[1:]
+    elif model_type == "Sequential":
+        return model.layers
+
+class TestSpatialSvdLayerSplitandSVDPrunner:
+
+
+    @pytest.mark.parametrize("model_type", ["Sequential", "Functional"])
+    @pytest.mark.parametrize("rank", [1024, 512])
+    def test_split_layer(self, model_type, rank):
+        """
+        test the output after and before the split_module call
+        """
+        model = get_model(model_type)
+        orig_conv_op = _get_layers(model, model_type)[2]
+
+        org_conv_op_shape = orig_conv_op.output_shape
+
+        layer1 = Layer(orig_conv_op, orig_conv_op.name, output_shape=org_conv_op_shape)
+
+        split_conv_op1, split_conv_op2 = SpatialSvdModuleSplitter.split_module(model, layer=layer1, rank=rank)
+
+        split_conv_output = split_conv_op2.output_shape
+
+        assert org_conv_op_shape == split_conv_output
+
+        # check the bias value after split.
+
+        assert len(split_conv_op2.get_weights()) == len(orig_conv_op.get_weights())
+        if len(orig_conv_op.get_weights()) > 1:
+
+            orig_bias_out = orig_conv_op.get_weights()[1]
+            split_bias_out = split_conv_op2.get_weights()[1]
+
+            assert np.allclose(orig_bias_out, split_bias_out, atol=1e-4)
+
+        # First split conv op should not have bias
+        assert len(split_conv_op1.get_weights()) == 1
+
+    @pytest.mark.parametrize("model_type", ["Sequential", "Functional"])
+    def test_split_layer_with_stride(self, model_type):
+        """
+        test the conv2d split after and before split_module call with stride
+        """
+        model = get_model(model_type)
+        orig_conv_op = _get_layers(model, model_type)[1]
+
+        org_conv_op_shape = orig_conv_op.output_shape
+
+        layer1 = Layer(orig_conv_op, orig_conv_op.name, output_shape=org_conv_op_shape)
+
+        split_conv_op1, split_conv_op2 = SpatialSvdModuleSplitter.split_module(model, layer=layer1, rank=5)
+
+        split_conv_output = split_conv_op2.output_shape
+
+        assert org_conv_op_shape == split_conv_output
+
+        # check the bias value after split.
+
+        assert len(split_conv_op2.get_weights()) == len(orig_conv_op.get_weights())
+        if len(orig_conv_op.get_weights()) > 1:
+
+            orig_bias_out = orig_conv_op.get_weights()[1]
+            split_bias_out = split_conv_op2.get_weights()[1]
+
+            assert np.allclose(orig_bias_out, split_bias_out, atol=1e-4)
+
+        # First split conv op should not have bias
+        assert len(split_conv_op1.get_weights()) == 1
+
+
+    @pytest.mark.parametrize("model_type", ["Sequential", "Functional"])
+    def test_perform_svd_and_split_layer(self, model_type):
+
+        model = get_model(model_type)
+        layer_db = LayerDatabase(model)
+        layer = layer_db.find_layer_by_name(_get_layers(model, model_type)[2].name)
+        org_count = len(list(layer_db._compressible_layers.values()))
+        splitter = SpatialSvdPruner()
+        splitter._perform_svd_and_split_layer(layer, 1024, layer_db)
+
+        assert layer not in list(layer_db._compressible_layers.values())
+
+        after_split_count = len(list(layer_db._compressible_layers.values()))
+        assert (org_count + 1) == after_split_count
+