AutoQuant¶
This notebook shows a working code example of how to use AIMET AutoQuant feature.
AIMET offers a suite of neural network post-training quantization techniques. Often, applying these techniques in a specific sequence, results in better accuracy and performance. Without the AutoQuant feature, the AIMET user needs to manually try out various combinations of AIMET quantization features. This manual process is error-prone and often time-consuming.
The AutoQuant feature, analyzes the model, determines the sequence of AIMET quantization techniques and applies these techniques. In addition, the user can specify the amount of accuracy drop that can be tolerated, in the AutoQuant API. As soon as this threshold accuracy is reached, AutoQuant stops applying any additional quantization technique. In summary, the AutoQuant feature saves time and automates the quantization of the neural networks.
Overall flow¶
This notebook covers the following 1. Instantiate the example evaluation and training pipeline 2. Load a pretrained FP32 model 3. Prepare the model 4. Validate the model 5. Determine the baseline FP32 accuracy 6. Define constants and helper functions 7. Apply AutoQuant
What this notebook is not¶
This notebook is not designed to show state-of-the-art AutoQuant results. For example, it uses a relatively quantization-friendly model like Resnet18. Also, some optimization parameters are deliberately chosen to have the notebook execute more quickly.
Dataset¶
This notebook relies on the ImageNet dataset for the task of image classification. If you already have a version of the dataset readily available, please use that. Else, please download the dataset from appropriate location (e.g. https://image-net.org/challenges/LSVRC/2012/index.php#).
Note1: The ImageNet dataset typically has the following characteristics and the dataloader provided in this example notebook rely on these - Subfolders ‘train’ for the training samples and ‘val’ for the validation samples. Please see the pytorch dataset description for more details. - A subdirectory per class, and a file per each image sample
Note2: To speed up the execution of this notebook, you may use a reduced subset of the ImageNet dataset. E.g. the entire ILSVRC2012 dataset has 1000 classes, 1000 training samples per class and 50 validation samples per class. But for the purpose of running this notebook, you could perhaps reduce the dataset to say 2 samples per class. This exercise is left upto the reader and is not necessary.
Edit the cell below and specify the directory where the downloaded ImageNet dataset is saved.
[ ]:
DATASET_DIR = '/path/to/dataset' # Please replace this with a real directory
1. Example evaluation and training pipeline¶
The following is an example training and validation loop for this image classification task.
Does AIMET have any limitations on how the training, validation pipeline is written? Not really. We will see later that AIMET will modify the user’s model to create a QuantizationSim model which is still a PyTorch model. This QuantizationSim model can be used in place of the original model when doing inference or training.
Does AIMET put any limitation on the interface of the evaluate() or train() methods? Not really. You should be able to use your existing evaluate and train routines as-is.
[ ]:
import os
import torch
from Examples.common import image_net_config
from Examples.torch.utils.image_net_evaluator import ImageNetEvaluator
from Examples.torch.utils.image_net_trainer import ImageNetTrainer
from Examples.torch.utils.image_net_data_loader import ImageNetDataLoader
class ImageNetDataPipeline:
@staticmethod
def get_val_dataloader() -> torch.utils.data.DataLoader:
"""
Instantiates a validation dataloader for ImageNet dataset and returns it
"""
data_loader = ImageNetDataLoader(DATASET_DIR,
image_size=image_net_config.dataset['image_size'],
batch_size=image_net_config.evaluation['batch_size'],
is_training=False,
num_workers=image_net_config.evaluation['num_workers']).data_loader
return data_loader
@staticmethod
def evaluate(model: torch.nn.Module, use_cuda: bool) -> float:
"""
Given a torch model, evaluates its Top-1 accuracy on the dataset
:param model: the model to evaluate
:param use_cuda: whether or not the GPU should be used.
"""
evaluator = ImageNetEvaluator(DATASET_DIR, image_size=image_net_config.dataset['image_size'],
batch_size=image_net_config.evaluation['batch_size'],
num_workers=image_net_config.evaluation['num_workers'])
return evaluator.evaluate(model, iterations=None, use_cuda=use_cuda)
@staticmethod
def finetune(model: torch.nn.Module, epochs, learning_rate, learning_rate_schedule, use_cuda):
"""
Given a torch model, finetunes the model to improve its accuracy
:param model: the model to finetune
:param epochs: The number of epochs used during the finetuning step.
:param learning_rate: The learning rate used during the finetuning step.
:param learning_rate_schedule: The learning rate schedule used during the finetuning step.
:param use_cuda: whether or not the GPU should be used.
"""
trainer = ImageNetTrainer(DATASET_DIR, image_size=image_net_config.dataset['image_size'],
batch_size=image_net_config.train['batch_size'],
num_workers=image_net_config.train['num_workers'])
trainer.train(model, max_epochs=epochs, learning_rate=learning_rate,
learning_rate_schedule=learning_rate_schedule, use_cuda=use_cuda)
2. Load a pretrained FP32 model¶
For this example, we are going to load a pretrained resnet18 model from torchvision. Similarly, you can load any pretrained PyTorch model instead.
[ ]:
from torchvision.models import resnet18
model = resnet18(pretrained=True)
3. Prepare the Model¶
AIMET quantization features require the user’s model definition to follow certain guidelines. For example, functionals defined in forward pass should be changed to equivalent torch.nn.Module. AIMET user guide lists all these guidelines. The following ModelPreparer API uses new graph transformation feature available in PyTorch 1.9+ version and automates model definition changes required to comply with the above guidelines.
[ ]:
from aimet_torch.model_preparer import prepare_model
model = prepare_model(model)
4. Validate the model¶
AIMET provides a model validator utility to help check whether AIMET feature can be applied on a Pytorch model. The model validator currently checks for the following conditions: - No modules are reused - Operations have modules associated with them - Opeartions are not defined as Functionals
It is recommeneded to use the validate_model() API after using the prepare_model() API to identify any problem areas that were not automatically addressed by the model_preparer() API.
[ ]:
from aimet_torch.model_validator.model_validator import ModelValidator
input_shape = (1, 3, 224, 224)
result = ModelValidator.validate_model(model, model_input=torch.rand(input_shape))
5. Determine the baseline FP32 accuracy¶
Let’s determine the FP32 (floating point 32-bit) accuracy of this model using the evaluate() routine
We should decide whether to place the model on a CPU or CUDA device. This example code will use CUDA if available in your current execution environment. You can change this logic and force a device placement if needed.
[ ]:
use_cuda = False
if torch.cuda.is_available():
use_cuda = True
model.to(torch.device('cuda'))
[ ]:
accuracy = ImageNetDataPipeline.evaluate(model, use_cuda)
print(accuracy)
6. Define Constants and Helper functions¶
In this section the constants and helper functions needed to run this eaxmple are defined.
EVAL_DATASET_SIZE A typical value is 5000. To execute this example faster this value has been set to 50
CALIBRATION_DATASET_SIZE A typical value is 2000. To execute this example faster this value has been set to 20
BATCH_SIZE User sets the batch size. As an example, set to 10
The helper function **_create_sampled_data_loader()** returns a DataLoader based on the dataset and the number of samples provided.
[ ]:
import random
EVAL_DATASET_SIZE = 50
CALIBRATION_DATASET_SIZE = 20
BATCH_SIZE = 10
_subset_samplers = {}
def _create_sampled_data_loader(dataset, num_samples):
if num_samples not in _subset_samplers:
indices = random.sample(range(len(dataset)), num_samples)
_subset_samplers[num_samples] = SubsetRandomSampler(indices=indices)
return DataLoader(dataset,
sampler=_subset_samplers[num_samples],
batch_size=BATCH_SIZE)
Prepare unlabeled dataset¶
The AutoQuant feature utilizes an unlabeled dataset to achieve quantization. The class UnlabeledDatasetWrapper creates an unlabeled Dataset object from a labeled Dataset.
[ ]:
from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
class UnlabeledDatasetWrapper(Dataset):
def __init__(self, dataset):
self._dataset = dataset
def __len__(self):
return len(self._dataset)
def __getitem__(self, index):
images, _ = self._dataset[index]
return images
from torchvision import transforms
normalize = transforms.Normalize(mean=image_net_config.dataset['images_mean'],
std=image_net_config.dataset['images_std'])
image_size = image_net_config.dataset['image_size'],
val_transforms = transforms.Compose([
transforms.CenterCrop(image_size),
transforms.ToTensor(),
normalize])
from Examples.torch.utils.image_net_data_loader import ImageFolder
imagenet_dataset = ImageFolder(root=os.path.join(DATASET_DIR, 'val'), transform=val_transforms, num_samples_per_class=2)
unlabeled_imagenet_dataset = UnlabeledDatasetWrapper(imagenet_dataset)
unlabeled_imagenet_data_loader = _create_sampled_data_loader(unlabeled_imagenet_dataset, CALIBRATION_DATASET_SIZE)
Prepare the evaluation callback function¶
The eval_callback() function takes the model to evaluate and the number of samples to use as arguments. If the num_samples argument is None, the whole evaluation dataset is used to evaluate the model.
[ ]:
from typing import Optional
def eval_callback(model: torch.nn.Module, num_samples: Optional[int] = None) -> float:
return ImageNetDataPipeline.evaluate(model, use_cuda)
7. Apply AutoQuant¶
As a first step, the AutoQuant object is created.
The allowed_accuracy_drop parameter is set by the user to convey to the AutoQuant feature, how much accuracy drop is tolerated by the user. AutoQuant applies a series of quantization features. When the allowed accuracy is reached, AutoQuant stops applying any subsequent quantization feature. Please refer AutoQuant User Guide and API documentation for complete details.
[ ]:
from aimet_torch.auto_quant import AutoQuant
auto_quant = AutoQuant(allowed_accuracy_drop=0.01,
unlabeled_dataset_iterable=unlabeled_imagenet_data_loader,
eval_callback=eval_callback)
Optionally set AdaRound Parameters¶
The AutoQuant feature internally uses default parameters to execute the AdaRound step. If and only if necessary, the default AdaRound Parameters should be modified using the API shown below.
Note: To execute this example faster, the default value of the num_iterations parameter has been reduced from 10000 to 2000
[ ]:
from aimet_torch.adaround.adaround_weight import AdaroundParameters
ADAROUND_DATASET_SIZE = 20
adaround_data_loader = _create_sampled_data_loader(unlabeled_imagenet_dataset, ADAROUND_DATASET_SIZE)
print(len(adaround_data_loader))
adaround_params = AdaroundParameters(adaround_data_loader, num_batches=len(adaround_data_loader), default_num_iterations=2000)
auto_quant.set_adaround_params(adaround_params)
Run AutoQuant¶
This step applies the AutoQuant feature. The best possible quantized model, the associated eval_score and the path to the AdaRound encoding files are returned.
[ ]:
dummy_input = torch.randn(input_shape)
model.eval()
model, accuracy, encoding_path =\
auto_quant.apply(model.cuda(),
dummy_input_on_cpu=dummy_input.cpu(),
dummy_input_on_gpu=dummy_input.cuda())
[ ]:
print(accuracy)
print(encoding_path)
Summary¶
Hope this notebook was useful for you to understand how to use AIMET AutoQuant feature.
Few additional resources - Refer to the AIMET API docs to know more details of the APIs and parameters - Refer to the other example notebooks to understand how to use AIMET CLE and AdaRound features in a standalone fashion.