# export_on_train_end

Python API: `luxonis_train.callbacks.export_on_train_end`

Exports the model to ONNX when training ends.

## Classes

### ExportOnTrainEnd

Export the model to ONNX when training ends.

The callback exports the best checkpoint. The `preferred_checkpoint` parameter of
[NeedsCheckpoint](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/callbacks/needs_checkpoint.md)
selects the best main metric or the lowest validation loss.

The callback does not build an NN Archive and does not convert the model for a device. Use
[ConvertOnTrainEnd](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/callbacks/convert_on_train_end.md)
for all three steps. When `trainer.callbacks` lists an active
[ConvertOnTrainEnd](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/callbacks/convert_on_train_end.md),
the config deactivates this callback. A
[ConvertOnTrainEnd](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/callbacks/convert_on_train_end.md)
that `trainer.smart_cfg_auto_populate` adds does not deactivate it.

#### Methods

##### on_train_end

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

Export the best checkpoint to ONNX.

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.export](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/core/core.md).
That method writes `<name>.onnx` to `<run_save_dir>/export`, where `<name>` is `exporter.name` or `model.name`. For a model with
one input, it also writes the `modelconverter` config `<name>.yaml`. It uploads these files to the run when
`exporter.upload_to_run` is set. When no checkpoint exists, the hook logs a warning and exports nothing.

The export loads the checkpoint into `pl_module` only for the export. After the hook, `pl_module` holds its earlier weights again.

Parameters

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