confusion_matrix

confusion_matrix

Confusion matrix logging callback for phase classification models.

ConfusionMatrixCallback

ConfusionMatrixCallback(log_every_n_epochs=1)

Bases: Callback

Callback to log confusion matrices for phase classification models.

Automatically detects when phase predictions are available in the metrics and logs confusion matrices to WandB.

Initialize confusion matrix callback.

Parameters:

Name	Type	Description	Default
log_every_n_epochs	int	Log confusion matrix every N epochs	1

Source code in src/versatil/training/callbacks/confusion_matrix.py

def __init__(self, log_every_n_epochs: int = 1):
    """Initialize confusion matrix callback.

    Args:
        log_every_n_epochs: Log confusion matrix every N epochs
    """
    super().__init__()
    self.log_every_n_epochs = log_every_n_epochs

on_train_epoch_end

on_train_epoch_end(trainer, pl_module)

Log training confusion matrix at end of epoch.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/confusion_matrix.py

def on_train_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Log training confusion matrix at end of epoch.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    if trainer.current_epoch % self.log_every_n_epochs != 0:
        return

    cm = pl_module.train_metrics.compute_confusion_matrix()
    if cm is not None:
        fig = self._create_confusion_matrix_figure(
            cm, "Train Phase Confusion Matrix"
        )
        if trainer.logger is not None:
            wandb_image = figure_to_wandb_image(fig)
            trainer.logger.log_metrics(
                {"train_phase_confusion_matrix": wandb_image},
                step=trainer.current_epoch,
            )
        plt.close(fig)

on_validation_epoch_end

on_validation_epoch_end(trainer, pl_module)

Log validation confusion matrix at end of epoch.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/confusion_matrix.py

def on_validation_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Log validation confusion matrix at end of epoch.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    if trainer.sanity_checking:
        return
    if trainer.current_epoch % self.log_every_n_epochs != 0:
        return
    cm = pl_module.val_metrics.compute_confusion_matrix()
    if cm is not None:
        fig = self._create_confusion_matrix_figure(cm, "Val Phase Confusion Matrix")
        if trainer.logger is not None:
            wandb_image = figure_to_wandb_image(fig)
            trainer.logger.log_metrics(
                {"val_phase_confusion_matrix": wandb_image},
                step=trainer.current_epoch,
            )
        plt.close(fig)