anomalib.utils.callbacks.nncf.callback

Callbacks for NNCF optimization.

Module Contents

Classes

NNCFCallback

Callback for NNCF compression.
class anomalib.utils.callbacks.nncf.callback.NNCFCallback(config: Dict, export_dir: str = None)[source]

Bases: pytorch_lightning.Callback

Callback for NNCF compression.

Assumes that the pl module contains a ‘model’ attribute, which is the PyTorch module that must be compressed.

Parameters

  • config (Dict) – NNCF Configuration

  • export_dir (Str) – Path where the export onnx and the OpenVINO xml and bin IR are saved. If None model will not be exported.

setup(trainer: pytorch_lightning.Trainer, pl_module: pytorch_lightning.LightningModule, stage: Optional[str] = None) → None[source]

Call when fit or test begins.

Takes the pytorch model and wraps it using the compression controller so that it is ready for nncf fine-tuning.

on_train_batch_start(trainer: pytorch_lightning.Trainer, _pl_module: pytorch_lightning.LightningModule, _batch: Any, _batch_idx: int, _unused: Optional[int] = 0) → None[source]

Call when the train batch begins.

Prepare compression method to continue training the model in the next step.

on_train_epoch_start(_trainer: pytorch_lightning.Trainer, _pl_module: pytorch_lightning.LightningModule) → None[source]

Call when the train epoch starts.

Prepare compression method to continue training the model in the next epoch.

on_train_end(_trainer: pytorch_lightning.Trainer, _pl_module: pytorch_lightning.LightningModule) → None[source]

Call when the train ends.

Exports onnx model and if compression controller is not None, uses the onnx model to generate the OpenVINO IR.