Post-Training Quantization with AIDGE#

What is neural network Quantization ?#

Deploying large neural network architectures on embedded targets can be a difficult task as they often require billions of floating point operations per inference.

To address this problem, several techniques have been developed over the past decades in order to reduce the computational load and energy consumption of those inferences. Those techniques include Pruning, Compression, Quantization and Distillation.

In particular, Post Training Quantization (PTQ) consists of taking an already trained network and replacing the costly floating-point MADD by their integer counterparts. The use of Bytes instead of Floats also leads to a smaller memory bandwidth.

alt text

While this process can seem trivial, the naive approach consisting only of rounding the parameters and activations doesn’t work in practice. Instead, we want to normalize the model in order to optimize the ranges of parameters and values propagated inside the network, before applying quantization.

The Quantization Pipeline#

The PTQ algorithm can be regarded as a processing pipeline:

We insert the scaling nodes that perform the calibration, called Quantizers and Dequantizers, inside the network.
We compute the value ranges at each processing node input and output, and optionally adjust them using a clipping policy.
We deduce from the ranges the scaling-factors and zero-points that must be stored inside the Quantizers.
Eventually, we insert casting nodes in the Quantizers to ensure that the outer dataflow uses integer precision.

Note that although the computational nodes are now in integer precision, the scaling-factors inside the Quantizers are still in floating-point numbers. To alleviate this, we propose an optional step, called Single-Shift Approximation (SSA), that replaces each floating point multiplication by an integer one followed by a Bit-Shift.

All this process can seem tedious, but as we will see later in this tutorial, all the required steps can be executed using a single call of our quantization API. We will also explain how to use the post quantization analysis API, which we refer to as the Diagnostic.

Overall, this notebook will show you how to perform PTQ of a convolutional neural network trained on the MNIST dataset. Let’s begin !

Environment setup#

Ensure that the Aidge modules are properly installed in the current environment. If it is the case, the following setup step can be skipped.

Note: When running this notebook on Binder, all required components are pre-installed.

[ ]:

%pip install aidge-core \
    aidge-onnx \
    aidge-backend-cpu \
    aidge-quantization

Import modules#

Besides Aidge modules, we need Numpy for manipulating the inputs, Matplotlib for visualization purposes, and gzip to uncompress the numpy dataset.

Then we want to import the Aidge modules:

the core module contains everything we need to manipulate the graph.
the backend module allows us to perform inferences using the CPU.
the onnx module allows us to load the pretrained model (stored in an onnx file).
the quantization module encaplsulate the Post Training Quantization API.

[ ]:

import gzip
import numpy as np
import matplotlib.pyplot as plt

import aidge_core
import aidge_onnx
import aidge_backend_cpu
import aidge_quantization
import aidge_model_explorer

print(" Available backends : ", aidge_core.Tensor.get_available_backends())

Exploring the original ONNX Model#

Download the ONNX file and Database (if needed)#

If git-lfs is not installed, the model and data can be downloaded using the following code snippet.

Reminder: This step is not needed when running in Binder.

[ ]:

BASE_URL = "https://gitlab.eclipse.org/eclipse/aidge/aidge/-/raw/main/examples/tutorials/PTQ_tutorial/"

# Download the model, input and output data files
files_to_download = ["ConvNet.onnx", "mnist_samples.npy.gz", "mnist_labels.npy.gz"]

for file_name in files_to_download:
    aidge_core.utils.download_file(
        file_path=file_name, file_url=f"{BASE_URL}{file_name}"
    )

Loading and Visualizing the Samples#

For the following we will use the provided samples of the MNIST database :

[ ]:

NB_SAMPLES = 1000

samples = np.load(gzip.GzipFile("mnist_samples.npy.gz", "r"))
labels = np.load(gzip.GzipFile("mnist_labels.npy.gz", "r"))

for i in range(10):
    plt.subplot(1, 10, i + 1)
    plt.axis("off")
    plt.tight_layout()
    plt.imshow(samples[i], cmap="gray")

Importing the model in Aidge#

The model we will quantize is a tiny Convolutional Network. Let’s load it and set it’s backend and datatype :

[ ]:

aidge_model = aidge_onnx.load_onnx("ConvNet.onnx")

aidge_core.remove_flatten(aidge_model)  # we want to get rid of the 'flatten' nodes ...

aidge_model.set_datatype(aidge_core.dtype.float32)
aidge_model.set_backend("cpu")

Performing example Inference#

We will now declare a helper function to perform the inferences. The function will prepare our input tensor, create a Scheduler, propagate the sample using it and output the class vector :

[ ]:

def propagate(model, sample):
    # setup the input
    sample = np.reshape(sample, (1, 1, 28, 28))
    input_tensor = aidge_core.Tensor(sample)
    input_tensor.to_dtype(aidge_core.dtype.float32)
    # create the Scheduler
    scheduler = aidge_core.SequentialScheduler(model)
    # perform the inference
    scheduler.forward(True, [input_tensor])
    # gather the outputs
    output_node = model.get_ordered_outputs()[0][0]
    output_tensor = output_node.get_operator().get_output(0)
    return np.array(output_tensor)


print("\n EXAMPLE INFERENCES :")
for i in range(10):
    output_array = propagate(aidge_model, samples[i])
    print(labels[i], " -> ", np.round(output_array, 2))

Computing the Model Accuracy#

To compare the model performance before and after the PTQ, let’s declare an accuracy computation helper function :

[ ]:

def compute_accuracy(model, samples, labels):
    acc = 0
    for i, x in enumerate(samples):
        y = propagate(model, x)
        if labels[i] == np.argmax(y):
            acc += 1
    return acc / len(samples)


accuracy = compute_accuracy(aidge_model, samples[0:NB_SAMPLES], labels)
print(f"\n MODEL ACCURACY : {accuracy * 100:.3f}%")

Model Visualization#

Let’s visualize the model before applying the PTQ. For this, we use aidge_model_explorer, which displays the model as a graph in the browser.

[ ]:

aidge_model_explorer.visualize(aidge_model, "model_before_ptq")

Calibration Dataset Creation#

In order to perform the PTQ, we need a set of samples so that we can evaluate the value ranges of each node inputs and outputsw. It is known as the calibration dataset. Let’s build it :

[ ]:

tensors = []
for sample in samples[0:NB_SAMPLES]:
    sample = np.reshape(sample, (1, 1, 28, 28)).astype(np.float32)
    tensor = aidge_core.Tensor(sample)
    tensors.append(tensor)

Fake or True Quantization ?#

With our API, it’s possible to perform two types of quantization: one said to be True and one said to be Fake.

Indeed, even though our goal is to represent the model’s weights and activations using fewer bits (e.g. 8-bit integers instead of 32-bit floats), we often want to explore freely aggressive or exotic configurations (e.g. 4-bit PTQ).

In that case it is enough to only simulate integers using floating point numbers that contains integers. This is known as Fake quantization.

On the other hand, one may want to be able to execute the quantized network using an Aidge backend that proposes the required integer kernels (e.g. the cpu backend using 8-bit inputs). This is known as True quantization.

We will start by demonstrating 8-bit true quantization.

Creating the PTQ Instance#

All the PTQ rely on an instance of a PTQ object that is constructed over an AIDGE graphview. This instance contains the quantization configurations for each node of the network. It also exposes the process() method that will run the PTQ pipeline.

Let’s declare a PTQ object for our model and print it’s configurations :

[ ]:

ptq = aidge_quantization.PTQ(
    aidge_model,
    activation_precision=aidge_core.dtype.int8,
    accumulation_precision=aidge_core.dtype.int32,
    inference_precision=aidge_core.dtype.float32,
)

print(ptq)

Among all of those configurations, we want to emphatize on some of them :

use_zero_point : allows for asymmetric quantization
channel_wise : allows the weights to have a scaling factor per channel
clipping : policy used to optimize the value ranges after the calibration

We will make a first try with the most simple approach, that is a symmetric and layer-wise PTQ without SSA :

[ ]:

ptq.use_zero_point = False
ptq.channel_wise = False
ptq.single_shift_approximation = False

We are now ready to perform the quantization over our calibration dataset. Note that the process() function does not override the input model, but instead return a new, quantized model. Let’s quantize our model!

[ ]:

q_model = ptq.process(tensors)

Let’s look at our quantized graphview, and observe that it now contains quantizers after every computational node :

[ ]:

aidge_model_explorer.visualize(q_model, "model_after_ptq")

Let’s now evaluate the post quantization accuracy :

[ ]:

accuracy = compute_accuracy(q_model, samples[0:NB_SAMPLES], labels)
print(f"\n POST QUANT ACCURACY : {accuracy * 100:.3f}%")

Aggressive Quantization Schemes#

We will now switch to fake quantization to attempt 3-bit PTQ over the same network, with this time zero-point and channel-wise activated :

[ ]:

ptq = aidge_quantization.PTQ(
    aidge_model,
    activation_precision=aidge_core.dtype.int3,
    accumulation_precision=aidge_core.dtype.int32,
    inference_precision=aidge_core.dtype.float32,
)

ptq.fake_quantize = True
ptq.use_zero_point = True
ptq.channel_wise = True

q_model = ptq.process(tensors)

accuracy = compute_accuracy(q_model, samples[0:NB_SAMPLES], labels)
print(f"\n POST QUANT ACCURACY : {accuracy * 100:.3f}%")

We see that we still maintain a decent accuracy ! Also note the warning in the log, it indicates that the ReLUs are now replaced by Clips operators. This is a consequence of the use of zero-point.

Post Quantization Analysis#

Because of the very aggressive scheme we just used (3-bit), we do have a significant accuracy degradation.

To interpret this drop, the Aidge PTQ API proposes a tool, the diagnostic, that allows to display metrics that quantify the degradations induced by the quantization. It is done in a fine-grain, per-node, manner.

This diagnostic is contained in the PTQ instance and proposes by default to compute the Signal-Noise Ratio (SNR). Let’s try it:

[ ]:

diag = ptq.get_diagnostic()
report = diag.compute_quantization_errors(
    ["SNR"], isolate_error=False, producer_error=False
)