Repository: quic/aimet (language: python)

Issue:
Add Keras Layer Output Utility



Discussion / hints:

Fix patch:
diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/layer_output_utils.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/layer_output_utils.py
new file mode 100644
index 0000000000..e742cb00f4
--- /dev/null
+++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/layer_output_utils.py
@@ -0,0 +1,153 @@
+# /usr/bin/env python3.8
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2023, 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-@@
+# =============================================================================
+
+""" This module contains utilities to capture and save intermediate layer-outputs of a model. """
+
+import os
+from typing import Union, List, Tuple
+import re
+from collections import OrderedDict
+import json
+import tensorflow as tf
+from keras.utils import tf_utils
+from aimet_tensorflow.keras.quantsim import QcQuantizeWrapper, QcQuantizableMultiHeadAttention
+from aimet_common.layer_output_utils import SaveInputOutput
+from aimet_common.utils import AimetLogger
+
+logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.LayerOutputs)
+
+class LayerOutputUtil:
+    """
+    This class captures output of every layer of a keras (fp32/quantsim) model, creates a layer-output name to
+    layer-output dictionary and saves the per layer outputs
+    """
+    def __init__(self, model: tf.keras.Model, save_dir: str = f"./KerasLayerOutput"):
+        """
+        Constructor - It initializes a few things that are required for capturing and naming layer-outputs.
+        :param model: Keras (fp32/quantsim) model.
+        :param save_dir: Directory to save the layer outputs.
+        """
+        # Freeze the model weights and state
+        model.trainable = False
+
+        # Get intermediate model for layer-outputs
+        self.intermediate_model = self._get_intermediate_model(model)
+
+        # Get actual Layer output name to modified layer output name dict
+        self.original_name_to_modified_name_mapper = self._get_original_name_to_modified_name_mapper(model)
+
+        # Saving the actual layer output name to modified layer output name (valid file name to save) in a json file
+        os.makedirs(save_dir, exist_ok=True)
+        with open(os.path.join(save_dir, "LayerOutputNameMapper.json"), 'w') as fp:
+            json.dump(self.original_name_to_modified_name_mapper, fp=fp, indent=4)
+
+        # Identify the axis-layout used for representing an image tensor
+        axis_layout = 'NHWC' if tf.keras.backend.image_data_format() == 'channels_last' else 'NCHW'
+
+        # Utility to save model inputs and their corresponding layer-outputs
+        self.save_inp_out_obj = SaveInputOutput(save_dir, axis_layout=axis_layout)
+
+    @classmethod
+    def _get_layer_output_name(cls, layer: Union[QcQuantizeWrapper, QcQuantizableMultiHeadAttention, tf.keras.layers.Layer]):
+        """
+        This function returns the actual layer output name for a given layer
+        :param layer: Keras model layer.
+        :return: Actual layer output name for the layer
+        """
+        if isinstance(layer, QcQuantizeWrapper):
+            return layer.original_layer.output.name
+        return layer.output.name
+
+    @classmethod
+    def _get_intermediate_model(cls, model: tf.keras.Model):
+        """
+        This function instantiates the feature extraction model for per layer outputs
+        :param model: Keras model.
+        :return: Intermediate keras model for feature extraction
+        """
+        outputs = [layer.output for layer in model.layers]
+        intermediate_model = tf.keras.models.Model(inputs=model.inputs, outputs=outputs)
+        intermediate_model.trainable = False
+        return intermediate_model
+
+    @classmethod
+    def _get_original_name_to_modified_name_mapper(cls, model: tf.keras.Model):
+        """
+        This function captures the per-layer output name and modifies it to make a valid file name
+        (by removing non-word characters) so that the layer output can be easily saved with the modified name.
+        :param model: Keras model.
+        :return: Actual layer name to modified layer name dict
+        """
+        original_name_to_modified_name_mapper = OrderedDict()
+        for layer in model.layers:
+            layer_output_name = cls._get_layer_output_name(layer)
+
+            # Replace all non-word characters with "_" to make it a valid file name for saving the results
+            # For Eg.: "conv2d/BiasAdd:0" gets converted to "conv2d_BiasAdd_0"
+            modified_layer_output_name = re.sub(r'\W+', "_", layer_output_name)
+
+            original_name_to_modified_name_mapper[layer_output_name] = modified_layer_output_name
+
+        return original_name_to_modified_name_mapper
+
+    def get_outputs(self, input_batch: Union[tf.Tensor, List[tf.Tensor], Tuple[tf.Tensor]]):
+        """
+        This function captures layer-outputs and renames them as per the AIMET exported model.
+        :param input_batch: Batch of inputs for which we want to obtain layer-outputs.
+        :return: layer-output name to layer-output batch dict
+        """
+        # Run in inference mode
+        outs = self.intermediate_model(input_batch, training=False)
+        output_pred = tf_utils.sync_to_numpy_or_python_type(outs)
+
+        return dict(zip(self.original_name_to_modified_name_mapper.values(), output_pred))
+
+    def generate_layer_outputs(self, input_batch: Union[tf.Tensor, List[tf.Tensor], Tuple[tf.Tensor]]):
+        """
+        This method captures output of every layer of a keras model & saves the inputs and corresponding layer-outputs to disk.
+        This allows layer-output comparison either between original fp32 model and quantization simulated model or quantization
+        simulated model and actually quantized model on-target to debug accuracy miss-match issues.
+
+        :param input_batch: Batch of Inputs for which layer output need to be generated
+        :return: None
+        """
+
+        batch_layer_name_to_layer_output = self.get_outputs(input_batch)
+        self.save_inp_out_obj.save(tf_utils.sync_to_numpy_or_python_type(input_batch), batch_layer_name_to_layer_output)
+
+        logger.info("Layer Outputs Saved")


Test patch:
diff --git a/TrainingExtensions/tensorflow/test/python/eager/test_keras_layer_outputs.py b/TrainingExtensions/tensorflow/test/python/eager/test_keras_layer_outputs.py
new file mode 100644
index 0000000000..c6008a7a01
--- /dev/null
+++ b/TrainingExtensions/tensorflow/test/python/eager/test_keras_layer_outputs.py
@@ -0,0 +1,252 @@
+# /usr/bin/env python3.8
+# -*- mode: python -*-
+# =============================================================================
+#  @@-COPYRIGHT-START-@@
+#
+#  Copyright (c) 2023, 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 re
+import json
+import shutil
+from glob import glob
+from datetime import datetime
+import os
+import numpy as np
+import progressbar
+import tensorflow as tf
+from keras.models import Sequential
+from keras.layers import BatchNormalization, Dense, Conv2D, Flatten, AvgPool2D, MaxPool2D
+from aimet_tensorflow.keras.quantsim import QuantizationSimModel, QcQuantizeWrapper
+from aimet_tensorflow.keras.model_preparer import prepare_model
+from aimet_tensorflow.keras.layer_output_utils import LayerOutputUtil
+
+class DummyDataLoader(tf.keras.utils.Sequence):
+
+    def __init__(self, batch_size, data_count):
+        self.batch_size = batch_size
+        self.data_count = data_count
+        self.data = [np.random.rand(16, 16, 3).astype(np.float32) for _ in range(data_count)]
+
+    def on_epoch_end(self):
+        # if anything needs to be done in between two epochs
+        pass
+
+    def __getitem__(self, index):
+        return tf.convert_to_tensor(self.data[index*self.batch_size: (index+1)*self.batch_size], dtype=np.float32), tf.convert_to_tensor([np.random.choice([0, 1]) for _ in range(self.batch_size)], dtype=np.int32)
+
+    def __len__(self):
+        return int(np.ceil(self.data_count / self.batch_size))
+
+def dummy_forward_pass(model, input_batch):
+    _ = model.predict(input_batch)
+
+def keras_model():
+    """ Function for returning a basic keras model """
+
+    model = Sequential([
+        Conv2D(8, (2, 2), input_shape=(16, 16, 3,)),
+        BatchNormalization(momentum=.3, epsilon=.65),
+        AvgPool2D(),
+        MaxPool2D(),
+        BatchNormalization(momentum=.4, epsilon=.25),
+        Conv2D(4, (2, 2), activation=tf.nn.tanh, kernel_regularizer=tf.keras.regularizers.l2(0.5)),
+        Flatten(),
+        Dense(2, activation='softmax')])
+    return model
+
+def get_quantsim_artifacts(base_model):
+    # Using model preparer
+    base_model = prepare_model(base_model)
+    dummy_input = np.random.rand(1, 16, 16, 3)
+
+    sim = QuantizationSimModel(
+        model=base_model,
+        quant_scheme='tf_enhanced',
+        rounding_mode="nearest",
+        default_output_bw=8,
+        default_param_bw=8,
+        in_place=False,
+        config_file=None
+    )
+
+    sim.trainable = False
+
+    sim.compute_encodings(dummy_forward_pass,
+                          forward_pass_callback_args=dummy_input
+                          )
+    return sim
+
+class TestLayerOutputUtil:
+    def test_generate_layer_output(self):
+        """ Test whether layer-output files are correctly saved and are getting correctly loaded """
+
+        # Load the baseline model
+        base_model = keras_model()
+
+        # Get the QuantSim artifacts
+        qs_obj = get_quantsim_artifacts(base_model)
+
+        # Temporary output path to store outputs temporarily
+        temp_folder_name = f"temp_keras_layer_output_{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}"
+        save_dir = os.path.join(os.getcwd(), temp_folder_name)
+
+        # Required params for the layer output generation function
+        data_points = 4
+        batch_size = 3
+
+        # Get the dataLoader
+        dataloader = DummyDataLoader(data_count=data_points, batch_size=batch_size)
+
+        layer_output_util_obj = LayerOutputUtil(model=qs_obj.model, save_dir=save_dir)
+        for batch_num, inp_batch in enumerate(dataloader):
+            batch_x, _ = inp_batch
+            layer_output_util_obj.generate_layer_outputs(input_batch=batch_x)
+
+        # Verify number of inputs
+        assert data_points == len(glob(save_dir+"/inputs/*.raw")) ## Check #Inputs
+
+        # Verify number of layer outputs
+        assert data_points == len(glob(save_dir+"/outputs/*")) ## Check #Outputs
+
+        # Getting the actual layer output names
+        actual_layer_output_names = list()
+        unmodified_actual_layer_output_names = list()
+        for each_layer in qs_obj.model.layers:
+            layer_output_name = each_layer.output.name
+            if isinstance(each_layer, QcQuantizeWrapper):
+                layer_output_name = each_layer.original_layer.output.name
+            unmodified_actual_layer_output_names.append(layer_output_name)
+            layer_output_name = re.sub(r'\W+', "_", layer_output_name)
+            actual_layer_output_names.append(layer_output_name)
+
+        # Getting the saved layer output names
+        saved_layer_output_list = list()
+        for each_layer_output_name in glob(save_dir+"/outputs/layer_outputs_0/*.raw"):
+            each_layer_output_name = each_layer_output_name.split("/")[-1][:-4]
+            saved_layer_output_list.append(each_layer_output_name)
+
+        # Verify layer-output names
+        np.testing.assert_array_equal(np.array(sorted(actual_layer_output_names)), np.array(sorted(saved_layer_output_list)))
+
+        # Verify final layer output for all data points
+        cnt = 0
+        n_iterations = np.ceil(data_points/batch_size)
+        with progressbar.ProgressBar(max_value=n_iterations) as progress_bar:
+            for batch_num, input_batch in enumerate(dataloader):
+                batch_x, _ = input_batch
+                for inp_batch in batch_x:
+                    actual_output = qs_obj.model.predict(np.expand_dims(inp_batch, axis=0))
+                    last_layer_name = qs_obj.model.layers[-1].original_layer.output.name
+                    last_layer_name = re.sub(r'\W+', "_", last_layer_name)
+                    last_layer_file_name = f"{save_dir}/outputs/layer_outputs_{cnt}/{last_layer_name}.raw"
+                    saved_last_layer_output = np.fromfile(last_layer_file_name, dtype=np.float32)
+                    np.testing.assert_array_equal(actual_output[0], saved_last_layer_output)
+                    cnt += 1
+
+                progress_bar.update(batch_num+1)
+                if (batch_num+1) >= n_iterations:
+                    break
+
+        # Test the layer output name mapper file and dict
+        saved_layer_output_name_mapper = json.load(open(temp_folder_name+"/LayerOutputNameMapper.json", "r"))
+
+        # Verify number of layers
+        np.testing.assert_array_equal(np.array(unmodified_actual_layer_output_names),
+                                      np.array(list(saved_layer_output_name_mapper.keys())))
+
+        # Verify modified layer name for each of the layers
+        for layer_idx, unmodified_actual_layer_name in enumerate(unmodified_actual_layer_output_names):
+
+            # Test saved file layer output
+            assert actual_layer_output_names[layer_idx] == \
+                   saved_layer_output_name_mapper[unmodified_actual_layer_name]
+
+            # Test dict layer output
+            assert actual_layer_output_names[layer_idx] == \
+                   layer_output_util_obj.original_name_to_modified_name_mapper[unmodified_actual_layer_name]
+
+        # Removing the temporary output that was created (if all tests are passed)
+        shutil.rmtree(save_dir)
+
+    def test_get_quantsim_outputs(self):
+        base_model = keras_model()
+        qs_obj = get_quantsim_artifacts(base_model)
+
+        qs_model_actual_output_names = list()
+
+        # Getting the actual layer output names
+        for each_layer in qs_obj.model.layers:
+            layer_output_name = each_layer.output.name
+            if isinstance(each_layer, QcQuantizeWrapper):
+                layer_output_name = re.sub(r'\W+', "_", each_layer.original_layer.output.name)
+            qs_model_actual_output_names.append(layer_output_name)
+
+        layer_output_obj = LayerOutputUtil(qs_obj.model)
+        batch_size = 3
+        data_count = 4
+        dataloader = DummyDataLoader(batch_size, data_count)
+        layer_output_dict = layer_output_obj.get_outputs(dataloader[0][0])
+
+        # Verify whether outputs are generated for all the layers
+        for each_actual_output_name in qs_model_actual_output_names:
+            assert each_actual_output_name in layer_output_dict.keys(), f"Output not generated for " \
+                                                                        f"{each_actual_output_name}"
+
+        # Verify the final layer output
+        actual_output = qs_obj.model.predict(dataloader[0][0])
+        np.testing.assert_array_equal(actual_output, layer_output_dict[qs_model_actual_output_names[-1]])
+
+    def test_get_original_model_outputs(self):
+        base_model = keras_model()
+
+        base_model_actual_output_names = list()
+
+        for each_layer in base_model.layers:
+            layer_output_name = re.sub(r'\W+', "_", each_layer.output.name)
+            base_model_actual_output_names.append(layer_output_name)
+
+        layer_output_obj = LayerOutputUtil(base_model)
+        batch_size = 3
+        data_count = 4
+        dataloader = DummyDataLoader(batch_size, data_count)
+        layer_output_dict = layer_output_obj.get_outputs(dataloader[0][0])
+
+        # Verify whether outputs are generated for all the layers
+        for each_actual_output_name in base_model_actual_output_names:
+            assert each_actual_output_name in layer_output_dict.keys(), f"Output not generated for " \
+                                                                        f"{each_actual_output_name}"
+
+        # Verify the final layer output
+        actual_output = base_model.predict(dataloader[0][0])
+        np.testing.assert_array_equal(actual_output, layer_output_dict[base_model_actual_output_names[-1]])