source.utils.custom_classes.LossCurveCallback

  1#!/usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3"""
  4PyTorch Lightning Callback for Loss and Accuracy Curve Visualization.
  5
  6This module provides a custom callback that tracks and visualizes training
  7and validation metrics during model training, saving plots and raw data to
  8disk.
  9"""
 10__docformat__ = "numpy"
 11
 12import os
 13import json
 14import matplotlib.pyplot as plt
 15from pytorch_lightning.callbacks import Callback
 16from source.utils import config as cfg
 17
 18
 19class LossCurveCallback(Callback):
 20    """
 21    PyTorch Lightning callback for tracking and plotting loss curves.
 22
 23    This callback monitors training loss, validation loss, and validation
 24    accuracy throughout the training process. At the end of training, it
 25    generates and saves visualization plots and raw metric data.
 26
 27    Attributes
 28    ----------
 29    save_dir : str
 30        Directory path where plots and metrics will be saved.
 31    train_losses : list of float
 32        Training loss values collected at the end of each training epoch.
 33    val_losses : list of float
 34        Validation loss values collected at the end of each validation epoch.
 35    val_accs : list of float
 36        Validation accuracy values collected at the end of each validation
 37        epoch.
 38
 39    Examples
 40    --------
 41    >>> from pytorch_lightning import Trainer
 42    >>> callback = LossCurveCallback(save_dir="./outputs/curves")
 43    >>> trainer = Trainer(callbacks=[callback])
 44    >>> trainer.fit(model, datamodule=data_module)
 45
 46    Notes
 47    -----
 48    The callback creates three output files in the save_dir:
 49    - loss_curve.png: Plot of training and validation losses
 50    - val_acc_curve.png: Plot of validation accuracy
 51    - metrics.json: Raw metric data in JSON format
 52    """
 53
 54    def __init__(self, save_dir=cfg.LOSS_CURVES_PATH):
 55        """
 56        Initialize the LossCurveCallback.
 57
 58        Parameters
 59        ----------
 60        save_dir : str, optional
 61            Directory path where plots and metrics will be saved
 62            (default is cfg.LOSS_CURVES_PATH).
 63
 64        Notes
 65        -----
 66        The save directory is created automatically if it does not exist.
 67        """
 68
 69        super().__init__()
 70        self.save_dir = save_dir
 71        os.makedirs(self.save_dir, exist_ok=True)
 72        self.train_losses = []
 73        self.train_accs = []
 74        self.val_losses = []
 75        self.val_accs = []
 76
 77    # ---------- Train loss per epoch ----------
 78    def on_train_epoch_end(self, trainer, pl_module):
 79        """
 80        Called at the end of each training epoch to collect training loss.
 81
 82        Parameters
 83        ----------
 84        trainer : pytorch_lightning.Trainer
 85            The PyTorch Lightning trainer instance.
 86        pl_module : pytorch_lightning.LightningModule
 87            The LightningModule being trained.
 88
 89        Notes
 90        -----
 91        Extracts the 'train_loss' metric from trainer.callback_metrics and
 92        appends it to the train_losses list.
 93        """
 94
 95        metrics = trainer.callback_metrics
 96        if "train_loss" in metrics:
 97            self.train_losses.append(metrics["train_loss"].item())
 98        if "train_acc" in metrics:
 99            self.train_accs.append(metrics["train_acc"].item())
100
101    # ---------- Val loss and acc per epoch ----------
102    def on_validation_epoch_end(self, trainer, pl_module):
103        """
104        Called at the end of each validation epoch to collect validation
105        metrics.
106
107        Parameters
108        ----------
109        trainer : pytorch_lightning.Trainer
110            The PyTorch Lightning trainer instance.
111        pl_module : pytorch_lightning.LightningModule
112            The LightningModule being validated.
113
114        Notes
115        -----
116        Extracts 'val_loss' and 'val_acc' metrics from trainer.callback_metrics
117        and appends them to their respective lists.
118        """
119
120        metrics = trainer.callback_metrics
121        if "val_loss" in metrics:
122            self.val_losses.append(metrics["val_loss"].item())
123        if "val_acc" in metrics:
124            self.val_accs.append(metrics["val_acc"].item())
125
126    def on_train_end(self, trainer, pl_module):
127        """
128        Called at the end of training to generate and save plots and metrics.
129
130        Parameters
131        ----------
132        trainer : pytorch_lightning.Trainer
133            The PyTorch Lightning trainer instance.
134        pl_module : pytorch_lightning.LightningModule
135            The trained LightningModule.
136
137        Notes
138        -----
139        This method performs three main tasks:
140        1. Generates and saves a loss curve plot (loss_curve.png) showing
141           training loss and validation loss over epochs.
142        2. Generates and saves a validation accuracy plot (val_acc_curve.png)
143           if validation accuracy was tracked.
144        3. Saves all raw metric data to a JSON file (metrics.json) for later
145           analysis or reproduction.
146
147        All output files are saved to the directory specified in save_dir.
148        """
149
150        # ---------- Save curves as a PNG ----------
151        plt.figure()
152        plt.plot(self.train_losses, label="Train Loss")
153        if len(self.val_losses) > 0:
154            plt.plot(self.val_losses, label="Val Loss")
155        plt.legend()
156        plt.title("Loss Curves")
157        plt.xlabel("Epochs")
158        plt.ylabel("Loss")
159        plt.savefig(os.path.join(self.save_dir, "loss_curve.png"))
160        plt.close()
161
162        plt.figure()
163        plt.plot(self.train_accs, label="Train Accuracy")
164        if len(self.val_accs) > 0:
165            plt.plot(self.val_accs, label="Val Accuracy")
166        plt.legend()
167        plt.title("Accuracy Curves")
168        plt.xlabel("Epochs")
169        plt.ylabel("Accuracy")
170        plt.savefig(os.path.join(self.save_dir, "val_acc_curve.png"))
171        plt.close()
172
173        # ---------- Save raw data ----------
174        data = {
175            "train_losses": self.train_losses,
176            "val_losses": self.val_losses,
177            "train_accs": self.train_accs,
178            "val_accs": self.val_accs,
179        }
180
181        with open(os.path.join(self.save_dir, "metrics.json"), "w") as f:
182            json.dump(data, f)