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)

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 – pytorch model

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

• 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.

Returns

None

aimet_torch.visualize_model.visualize_weight_ranges(model, visualization_url, selected_layers=None)

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 – pytorch model

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

• 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.

Returns

None

aimet_torch.visualize_model.visualize_changes_after_optimization(old_model, new_model, visualization_url, selected_layers=None)

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 – pytorch model before optimization

• new_model – pytorch model after optimization

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

• 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.

Returns

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)