# -*- 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, Dict, List, Tuple
from enum import Enum
import shutil
import re
import numpy as np
import onnx
import torch
from aimet_common.utils import AimetLogger
from aimet_common.layer_output_utils import SaveInputOutput, save_layer_output_names
from aimet_torch.quantsim import ExportableQuantModule, QuantizationSimModel
from aimet_torch import utils
from aimet_torch import torchscript_utils
from aimet_torch.onnx_utils import OnnxSaver, OnnxExportApiArgs
from aimet_torch.qc_quantize_recurrent import QcQuantizeRecurrent
from aimet_torch.v2.nn.base import BaseQuantizationMixin
logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.LayerOutputs)
[docs]class NamingScheme(Enum):
""" Enumeration of layer-output naming schemes. """
PYTORCH = 1
""" Names outputs according to exported pytorch model. Layer names are used. """
ONNX = 2
""" Names outputs according to exported onnx model. Layer output names are generally numeric. """
TORCHSCRIPT = 3
""" Names outputs according to exported torchscript model. Layer output names are generally numeric. """
[docs]class LayerOutputUtil:
""" Implementation to capture and save outputs of intermediate layers of a model (fp32/quantsim). """
def __init__(self, model: torch.nn.Module, dir_path: str, naming_scheme: NamingScheme = NamingScheme.PYTORCH,
dummy_input: Union[torch.Tensor, Tuple, List] = None, onnx_export_args: Union[OnnxExportApiArgs, Dict] = None):
"""
Constructor for LayerOutputUtil.
:param model: Model whose layer-outputs are needed.
:param dir_path: Directory wherein layer-outputs will be saved.
:param naming_scheme: Naming scheme to be followed to name layer-outputs. There are multiple schemes as per
the exported model (pytorch, onnx or torchscript). Refer the NamingScheme enum definition.
:param dummy_input: Dummy input to model. Required if naming_scheme is 'NamingScheme.ONNX' or 'NamingScheme.TORCHSCRIPT'.
:param onnx_export_args: Should be same as that passed to quantsim export API to have consistency between
layer-output names present in exported onnx model and generated layer-outputs. Required if naming_scheme is
'NamingScheme.ONNX'.
"""
# Utility to capture layer-outputs
self.layer_output = LayerOutput(model=model, naming_scheme=naming_scheme, dir_path=dir_path, dummy_input=dummy_input,
onnx_export_args=onnx_export_args)
# Utility to save model inputs and their corresponding layer-outputs
self.save_input_output = SaveInputOutput(dir_path=dir_path, axis_layout='NCHW')
[docs] def generate_layer_outputs(self, input_batch: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.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 we want to obtain layer-outputs.
:return: None
"""
input_instance_count = len(input_batch) if isinstance(input_batch, torch.Tensor) else len(input_batch[0])
logger.info("Generating layer-outputs for %d input instances", input_instance_count)
# Obtain layer-output name to output dictionary
layer_output_batch_dict = self.layer_output.get_outputs(input_batch)
# Place inputs and layer-outputs on CPU
input_batch = LayerOutputUtil._get_input_batch_in_numpy(input_batch)
layer_output_batch_dict = LayerOutputUtil._get_layer_output_batch_in_numpy(layer_output_batch_dict)
# Save inputs and layer-outputs
self.save_input_output.save(input_batch, layer_output_batch_dict)
logger.info('Successfully generated layer-outputs for %d input instances', input_instance_count)
@staticmethod
def _get_input_batch_in_numpy(input_batch: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]]) -> \
Union[np.ndarray, List[np.ndarray], Tuple[np.ndarray]]:
"""
Coverts the torch tensors into numpy arrays
:param input_batch: input batch with torch tensors
:return: input batch with numpy arrays
"""
if isinstance(input_batch, (List, Tuple)):
numpy_input_batch = []
for ith_input in input_batch:
numpy_input_batch.append(ith_input.cpu().numpy())
return numpy_input_batch
return input_batch.cpu().numpy()
@staticmethod
def _get_layer_output_batch_in_numpy(layer_output_dict: Dict[str, torch.Tensor]) -> Dict[str, np.ndarray]:
"""
Converts the torch tensors into numpy arrays
:param layer_output_dict: layer output dictionary with torch tensors
:return: layer output dictionary with numpy arrays
"""
layer_output_numpy_dict = {}
for output_name, output_tensor in layer_output_dict.items():
layer_output_numpy_dict[output_name] = output_tensor.cpu().numpy()
return layer_output_numpy_dict
class LayerOutput:
"""
This class creates a layer-output name to layer-output dictionary. The layer-output names are as per the AIMET exported
pytorch/onnx/torchscript model.
"""
def __init__(self, model: torch.nn.Module, dir_path: str, naming_scheme: NamingScheme = NamingScheme.PYTORCH,
dummy_input: Union[torch.Tensor, Tuple, List] = None, onnx_export_args: Union[OnnxExportApiArgs, Dict] = None):
"""
Constructor - It initializes few dictionaries that are required for capturing and naming layer-outputs.
:param model: Model whose layer-outputs are needed.
:param dir_path: Directory wherein layer-output names arranged in topological order will be saved. It will also
be used to temporarily save onnx/torchscript equivalent of the given model.
:param naming_scheme: Naming scheme to be followed to name layer-outputs. There are multiple schemes as per
the exported model (pytorch, onnx or torchscript). Refer the NamingScheme enum definition.
:param dummy_input: Dummy input to model (required if naming_scheme is 'onnx').
:param onnx_export_args: Should be same as that passed to quantsim export API to have consistency between
layer-output names present in exported onnx model and generated layer-outputs (required if naming_scheme is
'onnx').
"""
self.model = model
self.module_to_name_dict = utils.get_module_to_name_dict(model=model, prefix='')
# Check whether the given model is quantsim model
self.is_quantsim_model = any(isinstance(module, (ExportableQuantModule, QcQuantizeRecurrent)) for module in model.modules())
# Obtain layer-name to layer-output name mapping
self.layer_name_to_layer_output_dict = {}
self.layer_name_to_layer_output_name_dict = {}
if naming_scheme == NamingScheme.PYTORCH:
for name, module in model.named_modules():
if utils.is_leaf_module(module) or isinstance(module, BaseQuantizationMixin):
name = name.replace('._module_to_wrap', '')
self.layer_name_to_layer_output_name_dict[name] = name
else:
self.layer_name_to_layer_output_name_dict = LayerOutput.get_layer_name_to_layer_output_name_map(
self.model, naming_scheme, dummy_input, onnx_export_args, dir_path)
# Replace any delimiter in layer-output name string with underscore
for layer_name, output_name in self.layer_name_to_layer_output_name_dict.items():
self.layer_name_to_layer_output_name_dict[layer_name] = re.sub(r'\W+', "_", output_name)
# Save layer-output names which are in topological order of model graph. This order can be used while comparing layer-outputs.
layer_output_names = list(self.layer_name_to_layer_output_name_dict.values())
save_layer_output_names(layer_output_names, dir_path)
def get_outputs(self, input_batch: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]]) -> Dict[str, torch.Tensor]:
"""
This function captures layer-outputs and renames them as per the AIMET exported pytorch/onnx/torchscript model.
:param input_batch: Batch of inputs for which we want to obtain layer-outputs.
:return: layer-name to layer-output batch dict
"""
# Fetch outputs of all the layers
self.layer_name_to_layer_output_dict = {}
if self.is_quantsim_model:
# Apply record-output hook to QuantizeWrapper modules (one node above leaf node in model graph)
utils.run_hook_for_layers_with_given_input(self.model, input_batch, self.record_outputs,
module_type_for_attaching_hook=(ExportableQuantModule, QcQuantizeRecurrent),
leaf_node_only=False)
else:
# Apply record-output hook to Original modules (leaf node in model graph)
utils.run_hook_for_layers_with_given_input(self.model, input_batch, self.record_outputs, leaf_node_only=True)
# Rename outputs according to pytorch/onnx/torchscript model
layer_output_name_to_layer_output_dict = LayerOutput.rename_layer_outputs(self.layer_name_to_layer_output_dict,
self.layer_name_to_layer_output_name_dict)
return layer_output_name_to_layer_output_dict
def record_outputs(self, module: torch.nn.Module, _, output: torch.Tensor):
"""
Hook function to capture output of a layer.
:param module: Layer-module in consideration.
:param _: Placeholder for the input of the layer-module.
:param output: Output of the layer-module.
:return: None
"""
layer_name = self.module_to_name_dict[module]
if isinstance(output, torch.Tensor):
self.layer_name_to_layer_output_dict[layer_name] = output.clone()
else:
logger.info("Skipping constant scalar output of layer %s", layer_name)
@staticmethod
def rename_layer_outputs(layer_name_to_layer_output_dict: Dict[str, torch.Tensor],
layer_name_to_layer_output_name_dict: Dict[str, str]) -> Dict[str, torch.Tensor]:
"""
Rename layer-outputs based on the layer-name to layer-output name map
:param layer_name_to_layer_output_dict: Dict containing layer-outputs
:param layer_name_to_layer_output_name_dict: Dict containing layer-output names
:return: layer_output_name_to_layer_output_dict
"""
layer_names = list(layer_name_to_layer_output_dict.keys())
for layer_name in layer_names:
if layer_name in layer_name_to_layer_output_name_dict:
# Rename the layer-output by using layer-output name, instead of layer-name
layer_output_name = layer_name_to_layer_output_name_dict[layer_name]
layer_name_to_layer_output_dict[layer_output_name] = layer_name_to_layer_output_dict.pop(layer_name)
else:
# Delete the layer-output as it doesn't have a name
layer_name_to_layer_output_dict.pop(layer_name)
return layer_name_to_layer_output_dict
@staticmethod
def get_layer_name_to_layer_output_name_map(model, naming_scheme: NamingScheme, dummy_input: Union[torch.Tensor, Tuple, List],
onnx_export_args: Union[OnnxExportApiArgs, Dict], dir_path: str) -> Dict[str, str]:
"""
This function produces layer-name to layer-output name map w.r.t the AIMET exported onnx/torchscript model. If a
layer gets expanded into multiple layers in the exported model then the intermediate layers are ignored and
output-name of last layer is used.
:param model: model
:param naming_scheme: onnx/torchscript
:param dummy_input: dummy input that is used to construct onnx/torchscript model
:param onnx_export_args: OnnxExportApiArgs instance same as that passed to quantsim export API
:param dir_path: directory to temporarily save the constructed onnx/torchscrip model
:return: dictionary of layer-name to layer-output name
"""
# Restore original model by removing quantization wrappers if present.
original_model = QuantizationSimModel.get_original_model(model)
# Set path to store exported onnx/torchscript model.
LayerOutput._validate_dir_path(dir_path)
exported_model_dir = os.path.join(dir_path, 'exported_models')
os.makedirs(exported_model_dir, exist_ok=True)
# Get node to i/o tensor name map from the onnx/torchscript model
if naming_scheme == NamingScheme.ONNX:
exported_model_node_to_io_tensor_map = LayerOutput.get_onnx_node_to_io_tensor_map(
original_model, exported_model_dir, dummy_input, onnx_export_args)
else:
exported_model_node_to_io_tensor_map = LayerOutput.get_torchscript_node_to_io_tensor_map(
original_model, exported_model_dir, dummy_input)
layer_names_list = [name for name, module in original_model.named_modules() if utils.is_leaf_module(module)]
layers_missing_in_exported_model = []
layer_name_to_layer_output_name_map = {}
# Get mapping between layer names and layer-output names.
logger.info("Layer Name to Layer Output-name Mapping")
# pylint: disable=protected-access
for layer_name in layer_names_list:
if layer_name in exported_model_node_to_io_tensor_map:
# pylint: disable=protected-access, unused-variable
layer_output_names, intermediate_layer_output_names = QuantizationSimModel._get_layer_activation_tensors(
layer_name, exported_model_node_to_io_tensor_map)
layer_name_to_layer_output_name_map[layer_name] = layer_output_names[0]
logger.info("%s -> %s", layer_name, layer_output_names[0])
else:
layers_missing_in_exported_model.append(layer_name)
if layers_missing_in_exported_model:
logger.warning("The following layers were not found in the exported model:\n"
"%s\n"
"This can be due to below reason:\n"
"\t- The layer was not seen while exporting using the dummy input provided in sim.export(). "
"Ensure that the dummy input covers all layers.",
layers_missing_in_exported_model)
# Delete onnx/torchscript models
shutil.rmtree(exported_model_dir, ignore_errors=False, onerror=None)
return layer_name_to_layer_output_name_map
@staticmethod
def get_onnx_node_to_io_tensor_map(model: torch.nn.Module, exported_model_dir: str, dummy_input: Union[torch.Tensor, Tuple, List],
onnx_export_args: Union[OnnxExportApiArgs, Dict]) -> Dict[str, Dict]:
"""
This function constructs an onnx model equivalent to the give pytorch model and then generates node-name to i/o
tensor-name map.
:param model: pytorch model without quantization wrappers
:param exported_model_dir: directory to save onnx model
:param dummy_input: dummy input to be used for constructing onnx model
:param onnx_export_args: configurations to generate onnx model
:return: onnx_node_to_io_tensor_map
"""
LayerOutput._validate_dummy_input(dummy_input)
LayerOutput._validate_onnx_export_args(onnx_export_args)
onnx_path = os.path.join(exported_model_dir, 'model.onnx')
OnnxSaver.create_onnx_model_with_pytorch_layer_names(onnx_model_path=onnx_path, pytorch_model=model,
dummy_input=dummy_input, onnx_export_args=onnx_export_args)
onnx_model = onnx.load(onnx_path)
onnx_node_to_io_tensor_map, _ = OnnxSaver.get_onnx_node_to_io_tensor_names_map(onnx_model)
return onnx_node_to_io_tensor_map
@staticmethod
def get_torchscript_node_to_io_tensor_map(model: torch.nn.Module, exported_model_dir: str,
dummy_input: Union[torch.Tensor, Tuple, List]) -> Dict[str, Dict]:
"""
This function constructs a torchscript model equivalent to the give pytorch model and then generates node-name to i/o
tensor-name map.
:param model: pytorch model without quantization wrappers
:param exported_model_dir: directory to save onnx model
:param dummy_input: dummy input to be used for constructing onnx model
:return: torchscript_node_to_io_tensor_map
"""
LayerOutput._validate_dummy_input(dummy_input)
ts_path = os.path.join(exported_model_dir, 'model.torchscript.pth')
with utils.in_eval_mode(model), torch.no_grad():
torchscript_utils.create_torch_script_model(ts_path, model, dummy_input)
trace = torch.jit.load(ts_path)
torch_script_node_to_io_tensor_map, _ = \
torchscript_utils.get_node_to_io_tensor_names_map(model, trace, dummy_input)
return torch_script_node_to_io_tensor_map
@staticmethod
def _validate_dir_path(dir_path: str):
"""
Validate directory path in which onnx/torchscript models will be temporarily saved
:param dir_path: directory path
:return:
"""
if dir_path is None:
raise ValueError("Missing directory path to save onnx/torchscript models")
@staticmethod
def _validate_dummy_input(dummy_input: Union[torch.Tensor, Tuple, List]):
"""
Validates dummy input which is used to generate onnx/torchscript model
:param dummy_input: single input instance
:return:
"""
if not isinstance(dummy_input, (torch.Tensor, tuple, list)):
raise ValueError("Invalid dummy_input data-type")
@staticmethod
def _validate_onnx_export_args(onnx_export_args: Union[OnnxExportApiArgs, Dict]):
"""
Validates export arguments which are used to generate an onnx model
:param onnx_export_args: export arguments
:return:
"""
if onnx_export_args is None:
onnx_export_args = OnnxExportApiArgs()
if not isinstance(onnx_export_args, (OnnxExportApiArgs, dict)):
raise ValueError("Invalid onnx_export_args data-type")