Source code for aimet_tensorflow.keras.layer_output_utils

# -*- 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 numpy as np
import tensorflow as tf
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)

[docs]class LayerOutputUtil: """ Implementation to capture and save outputs of intermediate layers of a model (fp32/quantsim) """ def __init__(self, model: tf.keras.Model, save_dir: str = "./KerasLayerOutput"): """ Constructor for LayerOutputUtil. :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', encoding='utf-8') 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 = [out.numpy() for out in outs] return dict(zip(self.original_name_to_modified_name_mapper.values(), output_pred))
[docs] 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 model & saves the inputs and corresponding layer-outputs to disk. :param input_batch: Batch of Inputs for which layer output need to be generated :return: None """ layer_output_batch_dict = self.get_outputs(input_batch) # Skip constant scalar layer-outputs const_scalar_layer_name = [] for layer_name, layer_output in layer_output_batch_dict.items(): if not isinstance(layer_output, np.ndarray): const_scalar_layer_name.append(layer_name) for layer_name in const_scalar_layer_name: logger.info("Skipping constant scalar output of layer %s", layer_name) _ = layer_output_batch_dict.pop(layer_name) self.save_inp_out_obj.save(np.array(input_batch), layer_output_batch_dict) logger.info("Layer Outputs Saved")