# aimet_callback

Python API: `luxonis_train.callbacks.aimet_callback`

Runs AIMET quantization when training ends.

## Classes

### AIMETCallback

Quantize the model with AIMET when training ends.

The callback passes the best checkpoint to
[LuxonisModel.quantize](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/core.md).
That method runs post-training quantization and quantization-aware training. It writes the quantized ONNX model and its NN Archive
to `<run_save_dir>/aimet`. The `exporter.aimet` section of the config holds the options.

[LuxonisLightningModule.configure_callbacks](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/lightning/luxonis_lightning.md)
adds this callback when `exporter.aimet.active` is set. The config does not have to list it. With that flag set, an entry in
`trainer.callbacks` adds a second instance, and the two instances quantize the model twice.

The advanced techniques are slow. AdaRound alone can take from 40 minutes to several hours. The time depends on the size of the
dataset and of the model.

#### Methods

##### init

```python
def __init__(**kwargs):
```

Initialize the callback.

Parameters

 * `**kwargs`: Keyword arguments forwarded to
   [NeedsCheckpoint](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/callbacks/needs_checkpoint.md).
   The `preferred_checkpoint` value `"loss"` selects the loss checkpoint. Every other value selects the metric checkpoint. Any
   other key raises `TypeError`.

##### on_train_end

```python
def on_train_end(_: pl.Trainer, pl_module: lxt.LuxonisLightningModule):
```

Quantize the model from the best checkpoint.

Lightning calls this hook once when `trainer.fit` ends. The hook selects a checkpoint with
[NeedsCheckpoint.get_checkpoint](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/callbacks/needs_checkpoint.md)
and passes it to
[LuxonisModel.quantize](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/core.md)
with the options of `exporter.aimet`. When no checkpoint exists, the quantization still runs, on the current weights of
`pl_module.core.lightning_module`. The quantization works on a deep copy of that module, so the module keeps its weights.

Parameters

 * `_` (`pl.Trainer`): The trainer. Unused.
 * `pl_module` (`lxt.LuxonisLightningModule`): The model to quantize.