fruit_project.utils.early_stop

Classes

EarlyStopping

Module Contents

class fruit_project.utils.early_stop.EarlyStopping(patience: int, delta: float, path: str, name: str, cfg: omegaconf.DictConfig, run: wandb.sdk.wandb_run.Run)

patience

delta

path

name

cfg

run

best_metric: float | None = None

counter = 0

earlystop = False

saved_checkpoints: List[Tuple[float, pathlib.Path]] = []

__call__(val_metric: float, model: torch.nn.Module) → bool

Checks if early stopping criteria are met and saves the model if the metric improves.

Parameters:

• val_metric (float) – Validation metric to monitor.

• model (nn.Module) – PyTorch model to save.

Returns:

True if early stopping criteria are met, False otherwise.

Return type:

bool

save_model(model: torch.nn.Module, val_metric: float)

Saves the model checkpoint.

Parameters:

• model (nn.Module) – PyTorch model to save.

• val_metric (float) – Validation metric value used for naming the checkpoint file.

Returns:

None

cleanup_checkpoints()

Deletes all saved checkpoints except the best one.

Returns:

None

get_best_model(model: torch.nn.Module) → torch.nn.Module

Loads the best model checkpoint and sets the model to evaluation mode.

Parameters:

model (nn.Module) – PyTorch model to load the best checkpoint into.

Returns:

The model with the best checkpoint loaded.

Return type:

nn.Module