# Applying quantization to an MNIST model In this tutorial, we will be providing a basic introduction to quantizing a model with CoreAI-Opt. After the end of this tutorial, you should be familiar with the following: 1. [How to apply CoreAI-Opt’s Weight-Only quantization](#Weight-Only-Quantization) 2. [How to apply CoreAI-Opt’s Weight + Activation quantization](#Weight-and-Activation-Quantization) 3. [How to apply CoreAI-Opt’s Quantization-Aware Training](#Quantization-Aware-Training) 4. [How to export CoreAI-Opt quantized models to Core AI](#Export-to-Core-AI) **Table of Contents:** - [Setup](#Setup) - [MNIST Dataset download](#MNIST-Dataset-download) - [Model definition](#Model-definition) - [Training and Evaluation](#Training-and-Evaluation) - [Train unquantized model](#Train-unquantized-model) - [Weight Only Quantization](#Weight-Only-Quantization) - [Weight and Activation Quantization](#Weight-and-Activation-Quantization) - [Quantization Aware Training](#Quantization-Aware-Training) - [Export to Core AI](#Export-to-Core-AI) ## Setup We will be using a basic CNN model and train it on the MNIST dataset and observe its final accuracy. Once we train this CNN model, we will apply quantization to it using `coreai-opt` starting with *weight only quantization (data free)*, then moving to apply calibration data based *weight + activation quantization*, and then finally do *Quantization Aware Training (QAT)*. ### MNIST Dataset download Helper to download the MNIST dataset with standard normalization applied. ### Model definition A simple CNN with a single Conv2d → ReLU → MaxPool block, followed by Flatten and a Linear classifier. ### Training and Evaluation Standard PyTorch training loop and evaluation function that computes accuracy. The CNN model used for this tutorial contains a single Conv2d, ReLU, MaxPool2d, Flatten, and Linear layer. Here’s the structure: ### Train unquantized model Let’s train this model (unquantized) so we can get a baseline accuracy. We save the trained weights so we can reload them for each quantization experiment. ## Weight Only Quantization Weight-only quantization compresses model weights to a lower precision (e.g., INT8) while keeping activations in full precision (FP32). This is the simplest form of quantization — it requires **no calibration data** and can be applied directly to a trained model. For this tutorial, we will use the `INT8` dtype. Refer to the `QuantizationSpec` reference for all options. `QuantizerConfig` describes how to quantize each operation through three spec keys: - `op_state_spec`: state tensors of the op (weights, biases). - `op_input_spec`: input activations — tensors flowing into the op. - `op_output_spec`: output activations — tensors flowing out of the op. For weight-only quantization, only the `weight` state spec is populated; setting `op_input_spec` and `op_output_spec` to `None` tells the quantizer to leave activations in full precision. > Note: `coreai-opt` also offers `QuantizerConfig.presets.w8()` as a shortcut for the above config. After calling `prepare()`, the model now has FakeQuantize modules inserted for weight tensors. These simulate quantization during the forward pass. ## Weight and Activation Quantization Now, let’s try using weight + activation quantization through `coreai-opt`. Weight + Activation quantization compresses both the model’s weights **and** intermediate activation tensors to a lower precision. This provides a greater speedup than weight-only, but requires **calibration data** — representative inputs are run through the model so it can compute appropriate scale and zero-point values for the quantized activations. For this tutorial, we will use the `INT8` dtype for activations. Refer to the `QuantizationSpec` reference for all options. Two changes from the weight-only config: - `op_input_spec` and `op_output_spec` are now populated, so activations are quantized too. - The `"*"` key applies the same INT8 spec to every operation input/output. We now call `calibration_mode()` and feed it representative inputs to populate scale and zero-point value. The `calibration_mode()` context manager enables range observers (to collect activation statistics) while disabling fake quantization (so the forward pass is numerically identical to the unquantized model). After exiting the context, observers are frozen and fake quantization is re-enabled. ## Quantization Aware Training Now, let’s try to fine-tune the weight + activation quantized model with QAT to recover any accuracy lost from quantization. The QAT training loop wraps each epoch in `wa_quantizer.training_mode()`: - During the context: fake quantization is enabled and observers are enabled (collecting activation statistics during training batches). - On exit: observers are frozen; fake quantization stays enabled, so the model is ready for evaluation with quant-aware weights. ## Export to Core AI Once the quantized model is ready, call `finalize()` to convert the fake quantization modules into the real quantized representation for deployment. Pass `ExportBackend.CoreAI` to `finalize(backend=...)` to target the `.aimodel` format produced by `coreai-torch`. We’ll export the finetuned weight + activation model from the previous section. The export proceeds in three steps: - Trace the model with `torch.export.export()` to obtain a graph representation. - Apply `cast_to_16_bit_precision()` to cast remaining FP32 parameters to FP16 for optimal on-device performance. - Convert the exported program to Core AI format using `coreai-torch.TorchConverter`.