Source code for aimet_torch.visualize_model

# -*- mode: python -*-
# =============================================================================
#  @@-COPYRIGHT-START-@@
#
#  Copyright (c) 2019-2021, 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-@@
# =============================================================================

""" Top level API for visualizing a pytorch model. """
import os
from typing import List
import torch
from bokeh import plotting
from bokeh.layouts import column
from aimet_torch import plotting_utils
from aimet_torch.utils import get_layer_by_name


[docs]def visualize_changes_after_optimization( old_model: torch.nn.Module, new_model: torch.nn.Module, results_dir: str, selected_layers: List = None ) -> List[plotting.figure]: """ Visualizes changes before and after some optimization has been applied to a model. :param old_model: pytorch model before optimization :param new_model: pytorch model after optimization :param results_dir: Directory to save the Bokeh plots :param selected_layers: a list of layers a user can choose to have visualized. If selected layers is None, all Linear and Conv layers will be visualized. :return: A list of bokeh plots """ file_path = os.path.join(results_dir, 'visualize_changes_after_optimization.html') plotting.output_file(file_path) subplots = [] if selected_layers: for name, module in new_model.named_modules(): if name in selected_layers and hasattr(module, "weight"): old_model_module = get_layer_by_name(old_model, name) new_model_module = module subplots.append( plotting_utils.visualize_changes_after_optimization_single_layer( name, old_model_module, new_model_module ) ) else: for name, module in new_model.named_modules(): if hasattr(module, "weight") and\ isinstance(module, (torch.nn.modules.conv.Conv2d, torch.nn.modules.linear.Linear)): old_model_module = get_layer_by_name(old_model, name) new_model_module = module subplots.append( plotting_utils.visualize_changes_after_optimization_single_layer( name, old_model_module, new_model_module ) ) plotting.save(column(subplots)) return subplots
[docs]def visualize_weight_ranges( model: torch.nn.Module, results_dir: str, selected_layers: List = None ) -> List[plotting.figure]: """ Visualizes weight ranges for each layer through a scatter plot showing mean plotted against the standard deviation, the minimum plotted against the max, and a line plot with min, max, and mean for each output channel. :param model: pytorch model :param selected_layers: a list of layers a user can choose to have visualized. If selected layers is None, all Linear and Conv layers will be visualized. :param results_dir: Directory to save the Bokeh plots :return: A list of bokeh plots """ file_path = os.path.join(results_dir, 'visualize_weight_ranges.html') plotting.output_file(file_path) subplots = [] if selected_layers: for name, module in model.named_modules(): if name in selected_layers and hasattr(module, "weight"): subplots.append(plotting_utils.visualize_weight_ranges_single_layer(module, name)) else: for name, module in model.named_modules(): if hasattr(module, "weight") and\ isinstance(module, (torch.nn.modules.conv.Conv2d, torch.nn.modules.linear.Linear)): subplots.append(plotting_utils.visualize_weight_ranges_single_layer(module, name)) plotting.save(column(subplots)) return subplots
[docs]def visualize_relative_weight_ranges_to_identify_problematic_layers( model: torch.nn.Module, results_dir: str, selected_layers: List = None ) -> List[plotting.figure]: """ For each of the selected layers, publishes a line plot showing weight ranges for each layer, summary statistics for relative weight ranges, and a histogram showing weight ranges of output channels with respect to the minimum weight range. :param model: pytorch model :param results_dir: Directory to save the Bokeh plots :param selected_layers: a list of layers a user can choose to have visualized. If selected layers is None, all Linear and Conv layers will be visualized. :return: A list of bokeh plots """ file_path = os.path.join(results_dir, 'visualize_relative_weight_ranges_to_identify_problematic_layers.html') plotting.output_file(file_path) subplots = [] # layer name -> module weights data frame mapping if not selected_layers: for name, module in model.named_modules(): if hasattr(module, "weight") and\ isinstance(module, (torch.nn.modules.conv.Conv2d, torch.nn.modules.linear.Linear)): subplots.append(plotting_utils.visualize_relative_weight_ranges_single_layer(module, name)) else: for name, module in model.named_modules(): if hasattr(module, "weight") and\ isinstance(module, (torch.nn.modules.conv.Conv2d, torch.nn.modules.linear.Linear)) and\ name in selected_layers: subplots.append(plotting_utils.visualize_relative_weight_ranges_single_layer(module, name)) plotting.save(column(subplots)) return subplots