AIMET API documentation (3)

1.1.4 Model Validation Utility

AIMET provides a model validator utility to help check whether AIMET functions can be applied to Pytorch models. The model validator currently checks the following conditions:

No modules are reused
An operation has a module associated with it and is not defined as a function (does not include a set of known operations)

In this section, we introduce models that fail validation checks and show how to run the model validator and fix the model so that the validation checks pass.

Example 1: Model with reused modules

We start with the following model, which contains two relu modules that share the same module instance.

class ModelWithReusedNodes(torch.nn.Module):
    """ Model that reuses a relu module. Expects input of shape (1, 3, 32, 32) """

    def __init__(self):
        super(ModelWithReusedNodes, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 8, kernel_size=2, stride=2, padding=2, bias=False)
        self.bn1 = torch.nn.BatchNorm2d(8)
        self.relu1 = torch.nn.ReLU(inplace=True)
        self.linear = torch.nn.Linear(2592, 10)

    def forward(self, *inputs):
        x = self.conv1(inputs[0])
        x = self.relu1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = x.view(x.size(0), -1)
        x = self.linear(x)
        return x

Import model validator:

from aimet_torch.model_validator.model_validator import ModelValidator

Run the model validator on the model by passing in the model along with the model inputs:

def validate_example_model():

    # Load the model to validate
    model = ModelWithReusedNodes()

    # Output of ModelValidator.validate_model will be True if model is valid, False otherwise
    ModelValidator.validate_model(model, model_input=torch.rand(1, 3, 32, 32))

For every validation check run on the model, the logger prints:

Utils - INFO - Running validator check <function validate_for_reused_modules at 0x7f127685a598>

If the validation check finds any issues with the model, the log will contain information on how to resolve the model:

Utils - WARNING - The following modules are used more than once in the model: ['relu1']
AIMET features are not designed to work with reused modules. Please redefine your model to use distinct modules for
each instance.

Finally, at the end of validation, any failed checks will be logged:

Utils - INFO - The following validator checks failed:
Utils - INFO -     <function validate_for_reused_modules at 0x7f127685a598>

In this case, the validate_for_reused_modules check informs that the relu1 module is used multiple times in the model. We rewrite the model by defining a separate relu instance for each use:

class ModelWithoutReusedNodes(torch.nn.Module):
    """ Model that is fixed to not reuse modules. Expects input of shape (1, 3, 32, 32) """

    def __init__(self):
        super(ModelWithoutReusedNodes, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 8, kernel_size=2, stride=2, padding=2, bias=False)
        self.bn1 = torch.nn.BatchNorm2d(8)
        self.relu1 = torch.nn.ReLU(inplace=True)
        self.relu2 = torch.nn.ReLU(inplace=True)
        self.linear = torch.nn.Linear(2592, 10)

    def forward(self, *inputs):
        x = self.conv1(inputs[0])
        x = self.relu1(x)
        x = self.bn1(x)
        x = self.relu2(x)
        x = x.view(x.size(0), -1)
        x = self.linear(x)
        return x

Now, after rerunning the model validator, all checks pass:

Utils - INFO - Running validator check <function validate_for_reused_modules at 0x7ff577373598>
Utils - INFO - Running validator check <function validate_for_missing_modules at 0x7ff5703eff28>
Utils - INFO - All validation checks passed.

Example 2: Model with functionals

We start with the following model, which uses a torch linear function layer in the forward pass:

class ModelWithFunctionalLinear(torch.nn.Module):
    """ Model that uses a torch functional linear layer. Expects input of shape (1, 3, 32, 32) """

    def __init__(self):
        super(ModelWithFunctionalLinear, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 8, kernel_size=2, stride=2, padding=2, bias=False)
        self.bn1 = torch.nn.BatchNorm2d(8)
        self.relu1 = torch.nn.ReLU(inplace=True)
        self.relu2 = torch.nn.ReLU(inplace=True)

    def forward(self, *inputs):
        x = self.conv1(inputs[0])
        x = self.relu1(x)
        x = self.bn1(x)
        x = self.relu2(x)
        x = x.view(x.size(0), -1)
        x = F.linear(x, torch.randn(10, 2592))
        return x

Running the model validator shows that the validate_for_missing_modules check failed:

Utils - INFO - Running validator check <function validate_for_missing_modules at 0x7f9dd9bd90d0>
Utils - WARNING - Ops with missing modules: ['matmul_8']
This can be due to several reasons:
1. There is no mapping for the op in ConnectedGraph.op_type_map. Add a mapping for ConnectedGraph to recognize and
be able to map the op.
2. The op is defined as a functional in the forward function, instead of as a class module. Redefine the op as a
class module if possible. Else, check 3.
3. This op is one that cannot be defined as a class module, but has not been added to ConnectedGraph.functional_ops.
Add to continue.
Utils - INFO - The following validator checks failed:
Utils - INFO -      <function validate_for_missing_modules at 0x7f9dd9bd90d0>

Checking has identified matmul_8 as a missing pytorch module. In this case it's due to reason #2 in the log, where the layer has been defined as a function in the forward function. To solve this problem, we rewrite the model by defining the layers as modules.

class ModelWithoutFunctionalLinear(torch.nn.Module):
    """ Model that is fixed to use a linear module instead of functional. Expects input of shape (1, 3, 32, 32) """

    def __init__(self):
        super(ModelWithoutFunctionalLinear, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 8, kernel_size=2, stride=2, padding=2, bias=False)
        self.bn1 = torch.nn.BatchNorm2d(8)
        self.relu1 = torch.nn.ReLU(inplace=True)
        self.relu2 = torch.nn.ReLU(inplace=True)
        self.linear = torch.nn.Linear(2592, 10)
        with torch.no_grad():
            self.linear.weight = torch.nn.Parameter(torch.randn(10, 2592))

    def forward(self, *inputs):
        x = self.conv1(inputs[0])
        x = self.relu1(x)
        x = self.bn1(x)
        x = self.relu2(x)
        x = x.view(x.size(0), -1)
        x = self.linear(x)
        return x

1.1.5 AIMET PyTorch Quantitative Analyzer API

AIMET PyTorch Quant Analyzer analyzes PyTorch models and points out quantification-sensitive layers in the model. It examines the model's sensitivity to weight and activation quantification, performing per-layer sensitivity and MSE analyses. It also exports the minimum and maximum range of encoding for each layer and the statistical histogram for each layer.

1.1.5.1 Top-level API

class aimet_torch.quant_analyzer.QuantAnalyzer(model, dummy_input, forward_pass_callback, eval_callback, modules_to_ignore=None)[source]

QuantAnalyzer tool provides

Model sensitivity to weight and activation quantization
Sensitivity analysis per layer
Encoding per layer (min - max range)

Parameters :

model (Module) – FP32 model for analytical quantification.
dummy_input (Union[Tensor, Tuple]) – Pass input to the model.
forward_pass_callback (CallbackFunc) – Callback function for model calibration that simply runs a forward pass on the model to calculate the encoding (delta/offset). This callback function should use representative data and should be a subset of the entire training/validation dataset (~1000 images/samples).
eval_callback (CallbackFunc) – Callback function used for model evaluation to determine model performance. This callback function is expected to return a scalar value representing the model performance evaluated against the entire test/evaluation dataset.
modules_to_ignore (Optional[List[Module]]) – exclude certain modules from analysis.

analyze(quant_scheme=<QuantScheme.post_training_tf_enhanced: 2>, default_param_bw=8, default_output_bw=8, config_file=None, results_dir=‘./tmp/’)[source]

Analyze the quantification of the model and point out sensitive parts/hot spots of the model through execution

model sensitivity to quantification,
Perform per-layer sensitivity analysis by enabling and disabling quantitation wrappers,
Export the minimum-maximum range for each layer of encoding,
Export each layer statistical histogram (PDF) when the quantization scheme is TF enhancement,
MSE analysis per layer

parameter:

quant_scheme (QuantScheme) – Quantization scheme. Supported values are QuantScheme.post_training_tf or QuantScheme.post_training_tf_enhanced.
default_param_bw (int) – Default bitwidth (4-31) used to quantize layer parameters.
default_output_bw (int) – Default bitwidth (4-31) used for quantization layer input and output.
config_file (Optional[str]) – Path to the model quantizer configuration file.
results_dir (str) – Directory to save results.

enable_per_layer_mse_loss(unlabeled_dataset_iterable, num_batches)[source]

Enable per-layer MSE loss analysis.
parameter:

unlabeled_dataset_iterable (Union[DataLoader[+T_co], Collection[+T_co]]) – Iterate over a collection of unlabeled datasets (i.e. use len to iterate). The values produced by this iteration are expected to be passed directly to the model.
num_batches (int) – Number of batches. It is recommended to use approximately 256 samples/image, so if the data loader's batch size is 64, 4 batches will result in 256 samples/image.

1.1.5.2 Code examples

Required import

from typing import Any
import torch
from torchvision import models
from aimet_common.defs import QuantScheme
from aimet_torch.model_preparer import prepare_model
from aimet_torch.quant_analyzer import QuantAnalyzer, CallbackFunc

Prepare to pass callback forward

# NOTE: In the actual use cases, the users should implement this part to serve
#       their own goals if necessary.
def forward_pass_callback(model: torch.nn.Module, _: Any = None) -> None:
    """
    NOTE: This is intended to be the user-defined model calibration function.
    AIMET requires the above signature. So if the user's calibration function does not
    match this signature, please create a simple wrapper around this callback function.

    A callback function for model calibration that simply runs forward passes on the model to
    compute encoding (delta/offset). This callback function should use representative data and should
    be subset of entire train/validation dataset (~1000 images/samples).

    :param model: PyTorch model.
    :param _: Argument(s) of this callback function. Up to the user to determine the type of this parameter.
    E.g. could be simply an integer representing the number of data samples to use. Or could be a tuple of
    parameters or an object representing something more complex.
    """
    # User action required
    # User should create data loader/iterable using representative dataset and simply run
    # forward passes on the model.

Prepare eval callback

# NOTE: In the actual use cases, the users should implement this part to serve
#       their own goals if necessary.
def eval_callback(model: torch.nn.Module, _: Any = None) -> float:
    """
    NOTE: This is intended to be the user-defined model evaluation function.
    AIMET requires the above signature. So if the user's calibration function does not
    match this signature, please create a simple wrapper around this callback function.

    A callback function for model evaluation that determines model performance. This callback function is
    expected to return scalar value representing the model performance evaluated against entire
    test/evaluation dataset.

    :param model: PyTorch model.
    :param _: Argument(s) of this callback function. Up to the user to determine the type of this parameter.
    E.g. could be simply an integer representing the number of data samples to use. Or could be a tuple of
    parameters or an object representing something more complex.
    :return: Scalar value representing the model performance.
    """
    # User action required
    # User should create data loader/iterable using entire test/evaluation dataset, perform forward passes on
    # the model and return single scalar value representing the model performance.
    return .8

Prepare model

    model = models.resnet18(pretrained=True).cuda().eval()
    input_shape = (1, 3, 224, 224)
    dummy_input = torch.randn(*input_shape).cuda()
    prepared_model = prepare_model(model)

Create QuantAnalyzer object

    quant_analyzer = QuantAnalyzer(model=prepared_model,
                                   dummy_input=dummy_input,
                                   forward_pass_callback=forward_pass_callback_fn,
                                   eval_callback=eval_callback_fn)
    # Approximately 256 images/samples are recommended for MSE loss analysis. So, if the dataloader
    # has batch_size of 64, then 4 number of batches leads to 256 images/samples.
    quant_analyzer.enable_per_layer_mse_loss(unlabeled_dataset_iterable=unlabeled_data_loader, num_batches=4)

Run Quantitative Analyzer

    quant_analyzer.analyze(quant_scheme=QuantScheme.post_training_tf_enhanced,
                           default_param_bw=8,
                           default_output_bw=8,
                           config_file=None,
                           results_dir="./quant_analyzer_results/")