octopi 1.0__py3-none-any.whl

octopi/pytorch_lightning/optuna_pl_ddp.py

@@ -0,0 +1,273 @@
+import os
+import argparse
+import copick
+import torch
+from tqdm import tqdm
+from typing import Optional, Union, Tuple, List
+from collections import defaultdict
+import pytorch_lightning as pl
+import torch.distributed as dist
+from pytorch_lightning import Trainer
+from pytorch_lightning.loggers import MLFlowLogger
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from monai.data import DataLoader, Dataset, CacheDataset, decollate_batch
+from dotenv import load_dotenv
+from monai.transforms import (
+    Compose,
+    EnsureChannelFirstd,
+    Orientationd,
+    AsDiscrete,
+    RandFlipd,
+    RandRotate90d,
+    NormalizeIntensityd,
+    NormalizeIntensityd,
+    RandCropByLabelClassesd,    
+)
+from monai.networks.nets import UNet
+from monai.losses import TverskyLoss
+from monai.metrics import DiceMetric, ConfusionMatrixMetric
+import optuna
+from optuna.integration import PyTorchLightningPruningCallback
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description = "Hyperparamter tuning using PyTorch Lightning distributed data-parallel and Optuna."
+    )
+    parser.add_argument('--copick_config_path', type=str, default='copick_config_dataportal_10439.json')
+    parser.add_argument('--copick_user_name', type=str, default='user0')
+    parser.add_argument('--copick_segmentation_name', type=str, default='paintedPicks')
+    parser.add_argument('--train_batch_size', type=int, default=1)
+    parser.add_argument('--val_batch_size', type=int, default=1)
+    parser.add_argument('--num_random_samples_per_batch', type=int, default=16)
+    parser.add_argument('--learning_rate', type=float, default=1e-4)
+    parser.add_argument('--num_epochs', type=int, default=20)
+    parser.add_argument('--num_gpus', type=int, default=1)
+    parser.add_argument('--num_optuna_trials', type=int, default=10)
+    parser.add_argument('--pruning', action="store_true", help="Activate the pruning feature. `MedianPruner` stops unpromising trials at the early stages of training.")
+    return parser.parse_args()
+
+
+class Model(pl.LightningModule):
+    def __init__(
+        self, 
+        spatial_dims: int = 3,
+        in_channels: int = 1,
+        out_channels: int = 8,
+        channels: Union[Tuple[int, ...], List[int]] = (48, 64, 80, 80),
+        strides: Union[Tuple[int, ...], List[int]] = (2, 2, 1),
+        num_res_units: int = 1,
+        lr: float=1e-3):
+    
+        super().__init__()
+        self.save_hyperparameters()
+        self.model = UNet(
+            spatial_dims=self.hparams.spatial_dims,
+            in_channels=self.hparams.in_channels,
+            out_channels=self.hparams.out_channels,
+            channels=self.hparams.channels,
+            strides=self.hparams.strides,
+            num_res_units=self.hparams.num_res_units,
+        )
+        self.loss_fn = TverskyLoss(include_background=True, to_onehot_y=True, softmax=True)  # softmax=True for multiclass
+        self.metric_fn = DiceMetric(include_background=False, reduction="mean", ignore_empty=True)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch['image'], batch['label']
+        y_hat = self(x)
+        loss = self.loss_fn(y_hat, y)
+        return loss
+    
+    def validation_step(self, batch, batch_idx):
+        with torch.no_grad(): # This ensures that gradients are not stored in memory
+            x, y = batch['image'], batch['label']
+            y_hat = self(x)
+            metric_val_outputs = [AsDiscrete(argmax=True, to_onehot=self.hparams.out_channels)(i) for i in decollate_batch(y_hat)]
+            metric_val_labels = [AsDiscrete(to_onehot=self.hparams.out_channels)(i) for i in decollate_batch(y)]
+
+            # compute metric for current iteration
+            self.metric_fn(y_pred=metric_val_outputs, y=metric_val_labels)
+            metrics = self.metric_fn.aggregate(reduction="mean_batch")
+            for i,m in enumerate(metrics):
+                self.log(f"validation metric class {i+1}", m, prog_bar=True, on_epoch=True, sync_dist=True)
+            metric = torch.mean(metrics) # cannot log ndarray 
+            self.log('val_metric', metric, prog_bar=True, on_epoch=True, sync_dist=True) # sync_dist=True for distributed training
+        return {'val_metric': metric}
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
+    
+
+class CopickDataModule(pl.LightningDataModule):
+    def __init__(
+        self, 
+        copick_config_path: str, 
+        train_batch_size: int,
+        val_batch_size: int,
+        num_random_samples_per_batch: int):
+    
+        super().__init__()
+        self.train_batch_size = train_batch_size
+        self.val_batch_size = val_batch_size
+        
+        self.data_dicts, self.nclasses = self.data_from_copick(copick_config_path)
+        self.train_files = self.data_dicts[:int(len(self.data_dicts)//2)]
+        self.val_files = self.data_dicts[int(len(self.data_dicts)//2):]
+        print(f"Number of training samples: {len(self.train_files)}")
+        print(f"Number of validation samples: {len(self.val_files)}")
+
+        # Non-random transforms to be cached
+        self.non_random_transforms = Compose([
+            EnsureChannelFirstd(keys=["image", "label"], channel_dim="no_channel"),
+            NormalizeIntensityd(keys="image"),
+            Orientationd(keys=["image", "label"], axcodes="RAS")
+        ])
+
+        # Random transforms to be applied during training
+        self.random_transforms = Compose([
+            RandCropByLabelClassesd(
+                keys=["image", "label"],
+                label_key="label",
+                spatial_size=[96, 96, 96],
+                num_classes=self.nclasses,
+                num_samples=num_random_samples_per_batch 
+            ),
+            RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
+            RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),    
+        ])
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        self.train_ds = CacheDataset(data=self.train_files, transform=self.non_random_transforms, cache_rate=1.0)
+        self.train_ds = Dataset(data=self.train_ds, transform=self.random_transforms)
+        self.val_ds = CacheDataset(data=self.val_files, transform=self.non_random_transforms, cache_rate=1.0)
+        self.val_ds = Dataset(data=self.val_ds, transform=self.random_transforms)
+
+    def train_dataloader(self) -> DataLoader:
+        return DataLoader(
+                self.train_ds,
+                batch_size=self.train_batch_size,
+                shuffle=True,
+                num_workers=4,
+                persistent_workers=True,
+                pin_memory=torch.cuda.is_available(),
+            )
+
+    def val_dataloader(self) -> DataLoader:
+        return DataLoader(
+                self.val_ds,
+                batch_size=self.val_batch_size,
+                shuffle=False,  # Ensure the data order remains consistent
+                num_workers=4,
+                persistent_workers=True,
+                pin_memory=torch.cuda.is_available(),
+            )
+
+    @staticmethod
+    def data_from_copick(copick_config_path):
+        root = copick.from_file(copick_config_path)
+        nclasses = len(root.pickable_objects) + 1
+        data_dicts = []
+        target_objects = defaultdict(dict)
+        for object in root.pickable_objects:
+            if object.is_particle:
+                target_objects[object.name]['label'] = object.label
+                target_objects[object.name]['radius'] = object.radius
+        
+        data_dicts = []
+        for run in tqdm(root.runs[:8]):
+            tomogram = run.get_voxel_spacing(10).get_tomogram('wbp').numpy()
+            segmentation = run.get_segmentations(name='paintedPicks', user_id='user0', voxel_size=10, is_multilabel=True)[0].numpy()
+            membrane_seg = run.get_segmentations(name='membrane', user_id="data-portal")[0].numpy()
+            segmentation[membrane_seg==1] = 1  
+            data_dicts.append({"image": tomogram, "label": segmentation})
+        
+        return data_dicts, nclasses
+
+
+def objective(trial: optuna.trial.Trial) -> float:
+    args = get_args()
+    mlf_logger = MLFlowLogger(experiment_name='training-3D-UNet-model-for-the-cryoET-ML-Challenge',
+                              tracking_uri='http://mlflow.mlflow.svc.cluster.local:5000',
+                              #run_name='test1'
+                              )
+
+    # Trainer callbacks
+    checkpoint_callback = ModelCheckpoint(monitor='val_metric', save_top_k=1, mode='max')
+    lr_monitor = LearningRateMonitor(logging_interval='epoch')
+
+    # Detect distributed training environment
+    devices = list(range(args.num_gpus))
+
+    #channels, strides_pattern, num_res_units = sync_hyperparameters(trial)
+    # We optimize the number of layers, strides, and number of residual units
+    num_layers = trial.suggest_int("num_layers", 3, 5)
+    base_channel = trial.suggest_categorical("base_channel", [8, 16, 32, 64])
+    channels = [base_channel * (2 ** i) for i in range(num_layers)]
+    num_downsampling_layers = trial.suggest_int("num_downsampling_layers", 1, num_layers - 1)
+    strides_pattern = [2] * num_downsampling_layers + [1] * (num_layers - num_downsampling_layers - 1)        
+    num_res_units = trial.suggest_int("num_res_units", 1, 3)
+    
+    model = Model(channels=channels, strides=strides_pattern, num_res_units=num_res_units, lr=args.learning_rate)
+    datamodule = CopickDataModule(args.copick_config_path, args.train_batch_size, args.val_batch_size, args.num_random_samples_per_batch)
+    callback = PyTorchLightningPruningCallback(trial, monitor="val_metric")
+
+    # Priotize performace over precision
+    torch.set_float32_matmul_precision('medium') # or torch.set_float32_matmul_precision('high')    
+    
+    # Trainer for distributed training with DDP
+    trainer = Trainer(
+        max_epochs=args.num_epochs,
+        logger=mlf_logger,
+        callbacks=[checkpoint_callback, lr_monitor],
+        strategy="ddp_spawn",
+        accelerator="gpu",
+        devices=devices,
+        num_nodes=1, #int(os.environ.get("WORLD_SIZE", 1)) // args.num_gpus,
+        log_every_n_steps=1
+    )
+
+    hyperparameters = dict(op_num_layers=num_layers, op_base_channel=base_channel, op_num_downsampling_layers=num_downsampling_layers, op_num_res_units=num_res_units)
+    trainer.logger.log_hyperparams(hyperparameters)
+    trainer.fit(model, datamodule=datamodule)
+    callback.check_pruned()
+    return trainer.callback_metrics["val_metric"].item()
+
+
+if __name__ == "__main__":
+    args = get_args()
+    # MLflow setup
+    username = os.getenv('MLFLOW_TRACKING_USERNAME')
+    password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+    if not password or not username:
+        print("Password not found in environment, loading from .env file...")
+        load_dotenv()  # Loads environment variables from a .env file
+        username = os.getenv('MLFLOW_TRACKING_USERNAME')
+        password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+        
+    # Check again after loading .env file
+    if not password:
+        raise ValueError("Password is not set in environment variables or .env file!")
+    else:
+        print("Password loaded successfully")
+        os.environ['MLFLOW_TRACKING_USERNAME'] = username
+        os.environ['MLFLOW_TRACKING_PASSWORD'] = password    
+    
+    pruner: optuna.pruners.BasePruner = (
+        optuna.pruners.MedianPruner() if args.pruning else optuna.pruners.NopPruner()
+    )
+    storage = "sqlite:///example.db"
+    study = optuna.create_study(
+        study_name="pl_ddp",
+        storage=storage,
+        direction="maximize",
+        load_if_exists=True,
+        pruner=pruner
+    )
+    study.optimize(objective, n_trials=args.num_optuna_trials)
+
+    # Print the best hyperparameters
+    print(f"Best trial: {study.best_trial.value}")
+    print(f"Best hyperparameters: {study.best_trial.params}")

octopi/pytorch_lightning/train_pl.py

@@ -0,0 +1,244 @@
+import os
+import argparse
+from typing import Optional, Union, Tuple, List
+from collections import defaultdict
+import torch
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning import Trainer
+from pytorch_lightning.loggers import MLFlowLogger
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.strategies import DDPStrategy
+import os
+import copick
+from tqdm import tqdm
+from monai.data import DataLoader, Dataset, CacheDataset, decollate_batch
+from dotenv import load_dotenv
+from monai.transforms import (
+    Compose,
+    EnsureChannelFirstd,
+    Orientationd,
+    AsDiscrete,
+    RandFlipd,
+    RandRotate90d,
+    NormalizeIntensityd,
+    NormalizeIntensityd,
+    RandCropByLabelClassesd,    
+)
+from monai.networks.nets import UNet
+from monai.losses import TverskyLoss
+from monai.metrics import DiceMetric, ConfusionMatrixMetric
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description = "Train a 3d U-Net model with PyTorch Lightning supporting distributed training strategies."
+    )
+    parser.add_argument('--copick_config_path', type=str, default='copick_config_dataportal_10439.json')
+    parser.add_argument('--copick_user_name', type=str, default='user0')
+    parser.add_argument('--copick_segmentation_name', type=str, default='paintedPicks')
+    parser.add_argument('--train_batch_size', type=int, default=1)
+    parser.add_argument('--val_batch_size', type=int, default=1)
+    parser.add_argument('--num_random_samples_per_batch', type=int, default=16)
+    parser.add_argument('--learning_rate', type=float, default=1e-4)
+    parser.add_argument('--num_epochs', type=int, default=20)
+    parser.add_argument('--num_gpus', type=int, default=1)
+    return parser.parse_args()
+
+
+class Model(pl.LightningModule):
+    def __init__(
+        self, 
+        spatial_dims: int = 3,
+        in_channels: int = 1,
+        out_channels: int = 8,
+        channels: Union[Tuple[int, ...], List[int]] = (48, 64, 80, 80),
+        strides: Union[Tuple[int, ...], List[int]] = (2, 2, 1),
+        num_res_units: int = 1,
+        lr: float=1e-3):
+        
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.model = UNet(
+            spatial_dims=self.hparams.spatial_dims,
+            in_channels=self.hparams.in_channels,
+            out_channels=self.hparams.out_channels,
+            channels=self.hparams.channels,
+            strides=self.hparams.strides,
+            num_res_units=self.hparams.num_res_units,
+        )
+        self.loss_fn = TverskyLoss(include_background=True, to_onehot_y=True, softmax=True)  # softmax=True for multiclass
+        self.metric_fn = DiceMetric(include_background=False, reduction="mean", ignore_empty=True)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch['image'], batch['label']
+        y_hat = self(x)
+        loss = self.loss_fn(y_hat, y)
+        return loss
+    
+    def validation_step(self, batch, batch_idx):
+        x, y = batch['image'], batch['label']
+        y_hat = self(x)
+        metric_val_outputs = [AsDiscrete(argmax=True, to_onehot=self.hparams.out_channels)(i) for i in decollate_batch(y_hat)]
+        metric_val_labels = [AsDiscrete(to_onehot=self.hparams.out_channels)(i) for i in decollate_batch(y)]
+
+        # compute metric for current iteration
+        self.metric_fn(y_pred=metric_val_outputs, y=metric_val_labels)
+        metrics = self.metric_fn.aggregate(reduction="mean_batch")
+        for i,m in enumerate(metrics):
+            self.log(f"validation metric class {i+1}", m, prog_bar=True, on_epoch=True, sync_dist=True)
+        metric = torch.mean(metrics) # cannot log ndarray 
+        self.log('val_metric', metric, prog_bar=True, on_epoch=True, sync_dist=True) # sync_dist=True for distributed training
+        return {'val_metric': metric}
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
+
+
+class CopickDataModule(pl.LightningDataModule):
+    def __init__(
+        self, 
+        copick_config_path: str,
+        copick_user_name: str,
+        copick_segmentation_name: str, 
+        train_batch_size: int,
+        val_batch_size: int,
+        num_random_samples_per_batch: int):
+    
+        super().__init__()
+        self.train_batch_size = train_batch_size
+        self.val_batch_size = val_batch_size
+        
+        self.data_dicts, self.nclasses = self.data_from_copick(copick_config_path, copick_user_name, copick_segmentation_name)
+        self.train_files = self.data_dicts[:int(len(self.data_dicts)//2)]
+        self.val_files = self.data_dicts[int(len(self.data_dicts)//2):]
+        print(f"Number of training samples: {len(self.train_files)}")
+        print(f"Number of validation samples: {len(self.val_files)}")
+
+        # Non-random transforms to be cached
+        self.non_random_transforms = Compose([
+            EnsureChannelFirstd(keys=["image", "label"], channel_dim="no_channel"),
+            NormalizeIntensityd(keys="image"),
+            Orientationd(keys=["image", "label"], axcodes="RAS")
+        ])
+
+        # Random transforms to be applied during training
+        self.random_transforms = Compose([
+            RandCropByLabelClassesd(
+                keys=["image", "label"],
+                label_key="label",
+                spatial_size=[96, 96, 96],
+                num_classes=self.nclasses,
+                num_samples=num_random_samples_per_batch 
+            ),
+            RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
+            RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),    
+        ])
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        self.train_ds = CacheDataset(data=self.train_files, transform=self.non_random_transforms, cache_rate=1.0)
+        self.train_ds = Dataset(data=self.train_ds, transform=self.random_transforms)
+        self.val_ds = CacheDataset(data=self.val_files, transform=self.non_random_transforms, cache_rate=1.0)
+        self.val_ds = Dataset(data=self.val_ds, transform=self.random_transforms)
+
+    def train_dataloader(self) -> DataLoader:
+        return DataLoader(
+                self.train_ds,
+                batch_size=self.train_batch_size,
+                shuffle=True,
+                num_workers=4,
+                pin_memory=torch.cuda.is_available()
+            )
+
+    def val_dataloader(self) -> DataLoader:
+        return DataLoader(
+                self.val_ds,
+                batch_size=self.val_batch_size,
+                num_workers=4,
+                pin_memory=torch.cuda.is_available(),
+                shuffle=False,  # Ensure the data order remains consistent
+            )
+
+    @staticmethod
+    def data_from_copick(copick_config_path, copick_user_name, copick_segmentation_name):
+        root = copick.from_file(copick_config_path)
+        nclasses = len(root.pickable_objects) + 1
+        data_dicts = []
+        target_objects = defaultdict(dict)
+        for object in root.pickable_objects:
+            if object.is_particle:
+                target_objects[object.name]['label'] = object.label
+                target_objects[object.name]['radius'] = object.radius
+        
+        data_dicts = []
+        for run in tqdm(root.runs[:2]):
+            tomogram = run.get_voxel_spacing(10).get_tomogram('wbp').numpy()
+            segmentation = run.get_segmentations(name=copick_segmentation_name, user_id=copick_user_name, voxel_size=10, is_multilabel=True)[0].numpy()
+            membrane_seg = run.get_segmentations(name='membrane', user_id="data-portal")
+            if membrane_seg:
+                membrane_seg = run.get_segmentations(name='membrane', user_id="data-portal")[0].numpy()
+                segmentation[membrane_seg==1]=1    
+            data_dicts.append({"image": tomogram, "label": segmentation})
+        
+        return data_dicts, nclasses
+
+
+def train():
+    args = get_args()
+    mlf_logger = MLFlowLogger(experiment_name='training-3D-UNet-model-for-the-cryoET-ML-Challenge',
+                              tracking_uri='http://mlflow.mlflow.svc.cluster.local:5000',
+                              #run_name='test1'
+                              )
+    # Trainer callbacks
+    checkpoint_callback = ModelCheckpoint(monitor='val_metric', save_top_k=1, mode='min')
+    lr_monitor = LearningRateMonitor(logging_interval='epoch')
+
+    # Detect distributed training environment
+    devices = list(range(args.num_gpus))
+
+    # Initialize model
+    model = Model(lr=args.learning_rate)
+    datamodule = CopickDataModule(args.copick_config_path, args.copick_user_name, args.copick_segmentation_name,
+                                  args.train_batch_size, args.val_batch_size, args.num_random_samples_per_batch)
+
+    # Priotize performace over precision
+    torch.set_float32_matmul_precision('medium') # or torch.set_float32_matmul_precision('high')
+
+    # Trainer for distributed training with DDP
+    trainer = Trainer(
+        max_epochs=args.num_epochs,
+        logger=mlf_logger,
+        callbacks=[checkpoint_callback, lr_monitor],
+        strategy=DDPStrategy(find_unused_parameters=False),
+        accelerator="gpu",
+        devices=devices,
+        num_nodes=1, #int(os.environ.get("WORLD_SIZE", 1)) // args.num_gpus,
+        log_every_n_steps=1
+    )
+
+    trainer.fit(model, datamodule=datamodule)
+
+
+if __name__ == "__main__":
+    # MLflow setup
+    username = os.getenv('MLFLOW_TRACKING_USERNAME')
+    password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+    if not password or not username:
+        print("Password not found in environment, loading from .env file...")
+        load_dotenv()  # Loads environment variables from a .env file
+        username = os.getenv('MLFLOW_TRACKING_USERNAME')
+        password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+        
+    # Check again after loading .env file
+    if not password:
+        raise ValueError("Password is not set in environment variables or .env file!")
+    else:
+        print("Password loaded successfully")
+        os.environ['MLFLOW_TRACKING_USERNAME'] = username
+        os.environ['MLFLOW_TRACKING_PASSWORD'] = password
+    
+    train()

octopi/stopping_criteria.py

@@ -0,0 +1,143 @@
+import numpy as np
+
+class EarlyStoppingChecker:
+    """
+    A class to manage various early stopping criteria for model training.
+    """
+    
+    def __init__(self, 
+                 max_nan_epochs=15, 
+                 plateau_patience=20, 
+                 plateau_min_delta=0.001,
+                 stagnation_patience=50, 
+                 convergence_window=5, 
+                 convergence_threshold=0.005,
+                 val_interval=15,
+                 monitor_metric='avg_fbeta'):
+        """
+        Initialize early stopping parameters.
+        
+        Args:
+            max_nan_epochs: Maximum number of epochs with NaN loss before stopping
+            plateau_patience: Number of validation checks to wait for plateau detection
+            plateau_min_delta: Minimum change to qualify as improvement
+            stagnation_patience: Number of validation intervals to wait for best metric improvement
+            convergence_window: Window size for calculating improvement rate
+            convergence_threshold: Minimum improvement rate threshold
+            val_interval: Number of epochs between validation runs
+            monitor_metric: Primary metric to monitor for early stopping criteria
+        """
+        self.max_nan_epochs = max_nan_epochs
+        self.plateau_patience = plateau_patience
+        self.plateau_min_delta = plateau_min_delta
+        self.stagnation_patience = stagnation_patience
+        self.convergence_window = convergence_window
+        self.convergence_threshold = convergence_threshold
+        self.val_interval = val_interval
+        self.monitor_metric = monitor_metric
+        
+        # Counters
+        self.nan_counter = 0
+        
+        # Flags for detailed reporting
+        self.stopped_reason = None
+
+    def check_for_nan(self, epoch_loss):
+        """Check for NaN in the loss."""
+        if np.isnan(epoch_loss):
+            self.nan_counter += 1
+            if self.nan_counter > self.max_nan_epochs:
+                self.stopped_reason = f"NaN values in loss for more than {self.max_nan_epochs} epochs"
+                return True
+        else:
+            self.nan_counter = 0  # Reset the counter if loss is valid
+        return False
+
+    def check_for_plateau(self, results):
+        """Detect plateaus in validation metrics."""
+        if len(results[self.monitor_metric]) < self.plateau_patience + 1:
+            return False
+            
+        # Get the last 'patience' number of validation points
+        recent_values = [x[1] for x in results[self.monitor_metric][-self.plateau_patience:]]
+        # Find the max value in the window
+        max_value = max(recent_values)
+        # Find the min value in the window
+        min_value = min(recent_values)
+        
+        # If the range of values is small, consider it a plateau
+        if max_value - min_value < self.plateau_min_delta:
+            self.stopped_reason = f"{self.monitor_metric} plateaued for {self.plateau_patience} validations"
+            return True
+        
+        return False
+
+    def check_best_metric_stagnation(self, results):
+        """Stop if best metric hasn't improved for a number of validation intervals."""
+        if "best_metric_epoch" not in results or len(results[self.monitor_metric]) < self.stagnation_patience + 1:
+            return False
+            
+        # Get epoch of the best metric so far
+        best_epoch = results["best_metric_epoch"]
+        current_epoch = results[self.monitor_metric][-1][0]
+        
+        # Check if it's been more than 'patience' validation intervals
+        if (current_epoch - best_epoch) >= (self.stagnation_patience * self.val_interval):
+            self.stopped_reason = f"No improvement for {self.stagnation_patience} validation intervals"
+            return True
+            
+        return False
+
+    # def check_convergence_rate(self, results):
+    #     """Stop when improvement rate slows below threshold."""
+    #     if len(results[self.monitor_metric]) < self.convergence_window + 1:
+    #         return False
+        
+    #     # Calculate average improvement rate over window
+    #     recent_values = [x[1] for x in results[self.monitor_metric][-(self.convergence_window+1):]]
+    #     improvements = [recent_values[i+1] - recent_values[i] for i in range(self.convergence_window)]
+    #     avg_improvement = sum(improvements) / self.convergence_window
+        
+    #     if avg_improvement < self.convergence_threshold and avg_improvement > 0:
+    #         self.stopped_reason = f"Convergence rate ({avg_improvement:.6f}) below threshold"
+    #         return True
+            
+    #     return False
+
+    def should_stop_training(self, epoch_loss, results=None, check_metrics=False):
+        """
+        Comprehensive check for whether training should stop.
+        
+        Args:
+            epoch_loss: Current epoch's loss value
+            results: Dictionary containing training metrics history
+            check_metrics: Whether to also check validation metrics-based criteria
+            
+        Returns:
+            bool: True if training should stop, False otherwise
+        """
+        # Check for NaN in loss (can be done every epoch)
+        if self.check_for_nan(epoch_loss):
+            return True
+            
+        # Only check metric-based criteria if requested and results are provided
+        if check_metrics and results:
+            # Check for plateau in validation metrics
+            if self.check_for_plateau(results):
+                return True
+                
+            # Check if best metric hasn't improved for a while
+            if self.check_best_metric_stagnation(results):
+                return True
+                
+            # # Check if convergence rate has slowed down
+            # if self.check_convergence_rate(results):
+            #     return True
+                
+        return False
+    
+    def get_stopped_reason(self):
+        """Get the reason for stopping, if any."""
+        if self.stopped_reason:
+            return f"Early stopping triggered: {self.stopped_reason}"
+        return "No early stopping criteria met."