AIMET Visualization for Quantization API

Top-level API Quantization

aimet_torch.visualize_model.visualize_relative_weight_ranges_to_identify_problematic_layers(model, visualization_url, selected_layers=None, session_id='optimization', display=True)

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. Visualization_url is in the form: http://<host name>:<port number>/

Parameters

• model (Module) – pytorch model

• visualization_url (str) – user provided url with session id for the visualizations.

• selected_layers (Optional[List[~T]]) – 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 session_id : custom string descriptor to be associated with this session, appended to visualization url. :type display: bool :param display: a bool variable to indicate if output is to be displayed immediately. :return: None

aimet_torch.visualize_model.visualize_weight_ranges(model, visualization_url, selected_layers=None, session_id='optimization', display=True)

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. Visualization_url is in the form: http://<host name>:<port number>/

Parameters

• model (Module) – pytorch model

• visualization_url (str) – user inputted url with session id set as optimization for the visualizations.

• selected_layers (Optional[List[~T]]) – 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 session_id : custom string descriptor to be associated with this session, appended to visualization url. :type display: bool :param display: a bool variable to indicate if output is to be displayed immediately. :return: None

aimet_torch.visualize_model.visualize_changes_after_optimization(old_model, new_model, visualization_url, selected_layers=None, session_id='optimization', display=True)

Visualizes changes before and after some optimization has been applied to a model. Visualization_url is in the form: http://<host name>:<port number>/

Parameters

• old_model (Module) – pytorch model before optimization

• new_model (Module) – pytorch model after optimization

• visualization_url (str) – user inputted url with session id set as optimization for the visualizations.

• selected_layers (Optional[List[~T]]) – 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 session_id : custom string descriptor to be associated with this session, appended to visualization url. :type display: bool :param display: a flag to indicate if output is to be displayed immediately. :return: None

Code Examples

Required imports

import copy
import torch
from torchvision import models
from aimet_common.utils import start_bokeh_server_session
from aimet_torch.cross_layer_equalization import equalize_model
from aimet_torch.examples.imagenet_dataloader import ImageNetDataLoader
from aimet_torch.examples.supervised_classification_pipeline import \
    create_stand_alone_supervised_classification_evaluator
from aimet_torch.utils import IterFirstX
from aimet_torch import batch_norm_fold
from aimet_torch import visualize_model

Comparing Model After Optimization

def visualize_changes_in_model_after_and_before_cle():
    """
    Code example for visualizating model before and after Cross Layer Equalization optimization
    """
    visualization_url, process = start_bokeh_server_session(8002)
    model = models.resnet18(pretrained=True).to(torch.device('cpu'))
    model = model.eval()
    model_copy = copy.deepcopy(model)

    batch_norm_fold.fold_all_batch_norms(model_copy, (1, 3, 224, 224))

    equalize_model(model, (1, 3, 224, 224))
    visualize_model.visualize_changes_after_optimization(model_copy, model, visualization_url)

Visualizing weight ranges in Model

def visualize_weight_ranges_model():
    """
    Code example for model visualization
    """
    visualization_url, process = start_bokeh_server_session(8002)
    model = models.resnet18(pretrained=True).to(torch.device('cpu'))
    model = model.eval()

    batch_norm_fold.fold_all_batch_norms(model, (1, 3, 224, 224))

    # Usually it is observed that if we do BatchNorm fold the layer's weight range increases.
    # This helps in visualizing layer's weight
    visualize_model.visualize_weight_ranges(model, visualization_url)

Visualizing Relative weight ranges in Model

def visualize_relative_weight_ranges_model():
    """
    Code example for model visualization
    """
    visualization_url, process = start_bokeh_server_session(8002)
    model = models.resnet18(pretrained=True).to(torch.device('cpu'))
    model = model.eval()

    batch_norm_fold.fold_all_batch_norms(model, (1, 3, 224, 224))

    # Usually it is observed that if we do BatchNorm fold the layer's weight range increases.
    # This helps in finding layers which can be equalized to get better performance on hardware
    visualize_model.visualize_relative_weight_ranges_to_identify_problematic_layers(model, visualization_url)