anomalib.utils.callbacks.export

Callback that compresses a trained model by first exporting to .onnx format, and then converting to OpenVINO IR.

Module Contents

Classes

ExportCallback

Callback to compresses a trained model.

Attributes

logger

anomalib.utils.callbacks.export.logger[source]
class anomalib.utils.callbacks.export.ExportCallback(input_size: Tuple[int, int], dirpath: str, filename: str, export_mode: str)[source]

Bases: pytorch_lightning.Callback

Callback to compresses a trained model.

Model is first exported to .onnx format, and then converted to OpenVINO IR.

Parameters

  • input_size (Tuple[int, int]) – Tuple of image height, width

  • dirpath (str) – Path for model output

  • filename (str) – Name of output model

on_train_end(trainer, pl_module: anomalib.models.components.AnomalyModule) → None[source]

Call when the train ends.

Converts the model to onnx format and then calls OpenVINO’s model optimizer to get the .xml and .bin IR files.