# aimet_utils

Python API: `luxonis_train.core.utils.aimet_utils`

AIMET post-training quantization and quantization-aware training.

[LuxonisModel.quantize](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/core.md)
calls these helpers. AIMET is an optional dependency, installed with the `aimet` extra of `luxonis-train`. The module imports
`aimet_torch` only inside the functions, so the module itself imports without AIMET.

## Functions

### check_aimet_available

```python
def check_aimet_available():
```

Raise an error when the `aimet_torch` package is not installed.

The function looks for the package with `importlib.util.find_spec`. It does not import the package.

Raises

 * `ImportError`: When `aimet_torch` is not installed. The message tells the user to install `luxonis-train[aimet]`.

### get_ptq_calibration_loader

```python
def get_ptq_calibration_loader(val_dataset: torch_data.Dataset[LuxonisLoaderTorchOutput], collate_fn: Callable[[list[LuxonisLoaderTorchOutput]], Any], batch_size: int, num_workers: int, pin_memory: bool, max_calibration_images: int | None) -> DataLoader:
```

Build the loader that calibrates the post-training quantization.

[LuxonisModel.quantize](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/core.md)
passes the validation dataset itself, not
[LuxonisModel.val_loader](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/core.md).
So `trainer.n_validation_batches` does not limit the calibration.

With `max_calibration_images`, the loader reads only the first `max_calibration_images` samples of `val_dataset`. The function
logs an info message when this leaves out samples. When `val_dataset` has fewer samples than requested, the function logs a
warning and uses all samples. The loader keeps the order of the dataset and keeps the last incomplete batch.

> **Example**
> The example turns the logger off, so the info message does not show:

```pycon
>>> import torch
>>> from loguru import logger
>>> from torch.utils.data import TensorDataset, default_collate
>>> logger.disable("luxonis_train")
>>> loader = get_ptq_calibration_loader(
...     TensorDataset(torch.arange(10)),
...     collate_fn=default_collate,
...     batch_size=4,
...     num_workers=0,
...     pin_memory=False,
...     max_calibration_images=6,
... )
>>> logger.enable("luxonis_train")
>>> [batch.tolist() for (batch,) in loader]
[[0, 1, 2, 3], [4, 5]]
```

Parameters

 * `val_dataset` (`torch_data.Dataset[LuxonisLoaderTorchOutput]`): The dataset of the validation view. It must support `len` when
   `max_calibration_images` is not `None`.
 * `collate_fn` (`Callable[[list[LuxonisLoaderTorchOutput]], Any]`): The function that merges a list of samples into a batch, such
   as
   [BaseLoaderTorch.collate_fn](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/loaders/base_loader.md).
 * `batch_size` (`int`): The number of samples in a batch.
 * `num_workers` (`int`): The number of worker processes of the loader. `0` loads the samples in the main process.
 * `pin_memory` (`bool`): Copy the tensors of each batch into pinned memory before the loader returns them.
 * `max_calibration_images` (`int | None`): The maximum number of samples to read. `None` reads all samples.

Returns

 * `DataLoader`: The calibration loader.

### post_training_quantization

```python
def post_training_quantization(model: LuxonisLightningModule, dummy_inputs: Tensor, val_loader: DataLoader, save_dir: Path, quant_scheme: QuantScheme | None = None, default_output_bw: int = 8, default_param_bw: int = 8, default_data_type: QuantizationDataType | None = None, config_file: str | None = None, adaround: bool = False, adaround_iterations: int | None = None, adaround_reg_param: float = 0.01, adaround_beta_range: tuple[int, int] = (20, 2), adaround_warm_start: float = 0.2, fold_batch_norms: bool = False, cross_layer_equalization: bool = False, batch_norm_reestimation: bool = False, sequential_mse: bool = False) -> QuantizationSimModel:
```

Quantize a module with AIMET after training.

The function runs these steps:

 * It moves `model` and `dummy_inputs` to the GPU when CUDA is available, and puts `model` in eval mode.
 * With `fold_batch_norms` and without `batch_norm_reestimation`, it folds the batch norms of `model` into the preceding layers.
 * With `cross_layer_equalization`, it equalizes the weight ranges of consecutive layers of `model`.
 * With `adaround`, it learns the rounding of the weights on at most ⌈2000 ⁄ B⌉ batches of `val_loader`, where B is the batch
   size. AdaRound writes its files with the prefix `adaround` to `save_dir`. The next steps use the module that AdaRound returns
   instead of `model`.
 * It builds a `QuantizationSimModel` around the module with `in_place=True`.
 * With `sequential_mse`, it applies sequential MSE on `val_loader` with 20 candidates.
 * With `adaround`, it loads `adaround.encodings` from `save_dir` and freezes these parameter encodings.
 * It computes the encodings with a forward pass of the inputs of every batch of `val_loader`, and shows a progress bar.

The steps log an info message for the batch norm folding, the cross-layer equalization, and the sequential MSE.

Parameters

 * `model` (`LuxonisLightningModule`): The module to quantize.
 * `dummy_inputs` (`Tensor`): An input batch for the graph traces, such as a random tensor of shape `[1, C, H, W]`.
 * `val_loader` (`DataLoader`): The calibration loader, such as the result of
   [get_ptq_calibration_loader](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/utils/aimet_utils.md).
   Each batch is a pair of the inputs and the labels.
 * `save_dir` (`Path`): The directory for the AdaRound files.
 * `quant_scheme` (`QuantScheme | None`): The AIMET quantization scheme. `None` selects `QuantScheme.min_max`.
 * `default_output_bw` (`int`): The bit width of the activations.
 * `default_param_bw` (`int`): The bit width of the parameters.
 * `default_data_type` (`QuantizationDataType | None`): The data type of a quantized value. `None` selects
   `QuantizationDataType.int`.
 * `config_file` (`str | None`): The path of an AIMET config JSON file. `None` with `batch_norm_reestimation` selects the
   per-channel config of AIMET.
 * `adaround` (`bool`): Apply AdaRound.
 * `adaround_iterations` (`int | None`): The number of AdaRound iterations, passed to `AdaroundParameters`.
 * `adaround_reg_param` (`float`): The AdaRound regularization parameter.
 * `adaround_beta_range` (`tuple[int, int]`): The start and the end of the AdaRound beta annealing.
 * `adaround_warm_start` (`float`): The share of the AdaRound iterations during which the rounding loss has no effect.
 * `fold_batch_norms` (`bool`): Fold the batch norms before quantization. It has no effect with `batch_norm_reestimation`, because
   [quantization_aware_training](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/utils/aimet_utils.md)
   then folds them.
 * `cross_layer_equalization` (`bool`): Apply cross-layer equalization.
 * `batch_norm_reestimation` (`bool`): Whether
   [quantization_aware_training](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/utils/aimet_utils.md)
   re-estimates the batch norms. With `True`, the function skips the batch norm folding. It also selects the per-channel config of
   AIMET when `config_file` is `None`.
 * `sequential_mse` (`bool`): Apply sequential MSE.

Returns

 * `QuantizationSimModel`: The simulation, with the computed encodings. Its `model` is the quantized module.

Raises

 * `ImportError`: When `aimet_torch` is not installed.
 * `AssertionError`: When `val_loader` has no batch.
 * `ValueError`: When a parameter is not finite before the quantization, or after the batch norm folding, the cross-layer
   equalization, or AdaRound.

### quantization_aware_training

```python
def quantization_aware_training(sim: QuantizationSimModel, dummy_inputs: Tensor, train_loader: DataLoader, optimizer: Optimizer, scheduler: LRScheduler, epochs: int, fold_batch_norms: bool = False, batch_norm_reestimation: bool = False) -> LuxonisLightningModule:
```

Train the quantized module of an AIMET simulation.

The function trains `sim.model` in place, with the quantizers in the forward pass. It puts the module in training mode, and moves
it to the GPU when CUDA is available. It turns off the automatic optimization of the module for the loop, and restores the
previous value at the end, also after an error.

For each epoch, the function runs every batch of `train_loader` with a progress bar. For each batch, it computes the loss with
[LuxonisLightningModule.compute_training_loss](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/lightning/luxonis_lightning.md),
runs the backward pass, and steps `optimizer`. It steps `scheduler` once at the end of each epoch.

With `batch_norm_reestimation`, the function then re-estimates the batch norm statistics on `train_loader`. With
`fold_batch_norms` too, it folds the batch norms into the preceding layers. When AIMET cannot trace the graph of the quantized
module, the function logs a warning and skips the folding.

Parameters

 * `sim` (`QuantizationSimModel`): The simulation from
   [post_training_quantization](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/utils/aimet_utils.md).
 * `dummy_inputs` (`Tensor`): An input batch for the graph trace of the batch norm folding.
 * `train_loader` (`DataLoader`): The training loader. Each batch is a pair of the inputs and the labels.
 * `optimizer` (`Optimizer`): The optimizer of the parameters of `sim.model`.
 * `scheduler` (`LRScheduler`): The learning rate scheduler of `optimizer`.
 * `epochs` (`int`): The number of passes over `train_loader`.
 * `fold_batch_norms` (`bool`): Fold the batch norms after the re-estimation. It has no effect without `batch_norm_reestimation`.
 * `batch_norm_reestimation` (`bool`): Re-estimate the batch norm statistics after the training.

Returns

 * `LuxonisLightningModule`: `sim.model` after the training. The function does not put it back in eval mode.

Raises

 * `ImportError`: When `aimet_torch` is not installed.
 * `AssertionError`: When `train_loader` has no batch.