garmentiq.classification.train_pytorch_nn

  1"""Cross-validated training of a classification model from scratch."""
  2import torch
  3import torch.nn as nn
  4from torch.utils.data import DataLoader, Dataset
  5from typing import Callable, Type
  6from tqdm.auto import tqdm
  7import os
  8from sklearn.model_selection import StratifiedKFold
  9from garmentiq.utils.device import empty_cache
 10from garmentiq.classification.utils import (
 11    CachedDataset,
 12    seed_worker,
 13    train_epoch,
 14    validate_epoch,
 15    save_best_model,
 16    validate_train_param,
 17    validate_test_param,
 18)
 19
 20
 21def train_pytorch_nn(
 22    model_class: Type[torch.nn.Module],
 23    model_args: dict,
 24    dataset_class: Callable,
 25    dataset_args: dict,
 26    param: dict,
 27):
 28    """
 29    Trains a PyTorch neural network using k-fold cross-validation with early stopping and model checkpointing.
 30
 31    This function performs training and validation across multiple folds using stratified sampling.
 32    It manages model instantiation, training loops, early stopping, and saves the best model based on validation loss.
 33
 34    Args:
 35        model_class (Type[torch.nn.Module]): The class of the PyTorch model to instantiate.
 36                                            Must inherit from `torch.nn.Module`.
 37        model_args (dict): Dictionary of arguments used to initialize `model_class`.
 38        dataset_class (Callable): A callable class or function that returns a `torch.utils.data.Dataset`-compatible dataset.
 39        dataset_args (dict): Dictionary with dataset components:
 40            - 'metadata_df' (pandas.DataFrame): Metadata with labels, used for stratification.
 41            - 'raw_labels' (array-like): Raw class labels used by StratifiedKFold.
 42            - 'cached_images' (torch.Tensor): Preprocessed image tensor.
 43            - 'cached_labels' (torch.Tensor): Corresponding labels.
 44        param (dict): Dictionary of training hyperparameters and configuration values.
 45                      Required Keys:
 46                          - `optimizer_class` (type): PyTorch optimizer class (e.g., `torch.optim.Adam`).
 47                          - `optimizer_args` (dict): Arguments passed to the optimizer.
 48                      Optional Keys (with defaults and types):
 49                          - `device` (Union[str, torch.device]): Training device, e.g. `"cpu"`,
 50                            `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware acceleration is opt-in;
 51                            pass it explicitly to use a GPU or Apple Silicon. Default: `"cpu"`.
 52                          - `n_fold` (int): Number of stratified folds for cross-validation. Default: 5.
 53                          - `n_epoch` (int): Number of training epochs per fold. Default: 100.
 54                          - `patience` (int): Epochs to wait before early stopping. Default: 5.
 55                          - `batch_size` (int): Batch size for training and validation. Default: 64.
 56                          - `model_save_dir` (str): Directory to save model checkpoints. Default: `"./models"`.
 57                          - `seed` (int): Random seed for reproducibility. Default: 88.
 58                          - `seed_worker` (Callable): Function to seed workers in the DataLoader. Default: `seed_worker`.
 59                          - `max_workers` (int): Number of subprocesses for data loading. Default: `os.cpu_count()`.
 60                          - `best_model_name` (str): Filename for saving the best model. Default: `"best_model.pt"`.
 61
 62    Raises:
 63        ValueError: If any required key is missing from `param`.
 64        TypeError: If any parameter is of the wrong type.
 65        FileNotFoundError: If the model directory cannot be created or accessed.
 66
 67    Returns:
 68        None
 69    """
 70    # Validate and catch parameters
 71    validate_train_param(param)
 72
 73    # Prepare save directories
 74    os.makedirs(param["model_save_dir"], exist_ok=True)
 75    overall_best_loss = float("inf")
 76    best_model_path = os.path.join(param["model_save_dir"], param["best_model_name"])
 77
 78    kfold = StratifiedKFold(
 79        n_splits=param["n_fold"], shuffle=True, random_state=param["seed"]
 80    )
 81
 82    # Loop through each fold
 83    for fold, (train_idx, val_idx) in enumerate(
 84        kfold.split(dataset_args["metadata_df"], dataset_args["raw_labels"])
 85    ):
 86        print(f"\nFold {fold + 1}/{param['n_fold']}")
 87
 88        # Prepare datasets and dataloaders
 89        train_dataset = dataset_class(
 90            train_idx, dataset_args["cached_images"], dataset_args["cached_labels"]
 91        )
 92        val_dataset = dataset_class(
 93            val_idx, dataset_args["cached_images"], dataset_args["cached_labels"]
 94        )
 95
 96        g = torch.Generator()
 97        g.manual_seed(param["seed"])
 98
 99        train_loader = DataLoader(
100            train_dataset,
101            batch_size=param["batch_size"],
102            shuffle=True,
103            num_workers=param["max_workers"],
104            worker_init_fn=param["seed_worker"],
105            generator=g,
106            pin_memory=param["pin_memory"],
107            persistent_workers=param["persistent_workers"],
108        )
109        val_loader = DataLoader(
110            val_dataset,
111            batch_size=param["batch_size"],
112            shuffle=False,
113            num_workers=param["max_workers"],
114            worker_init_fn=param["seed_worker"],
115            generator=g,
116            pin_memory=param["pin_memory"],
117            persistent_workers=param["persistent_workers"],
118        )
119
120        # Initialize model and optimizer
121        model = model_class(**model_args).to(param["device"])
122        if param["device"].type == "cuda" and torch.cuda.device_count() > 1:
123            model = torch.nn.DataParallel(model)
124        optimizer = param["optimizer_class"](
125            model.parameters(), **param["optimizer_args"]
126        )
127        empty_cache(param["device"])
128
129        best_fold_loss = float("inf")
130        patience_counter = 0
131        epoch_pbar = tqdm(range(param["n_epoch"]), desc="Total Progress", leave=False)
132
133        # Training and Validation Loop
134        for epoch in epoch_pbar:
135            # Training phase
136            epoch_loss = train_epoch(model, train_loader, optimizer, param)
137            # Validation phase
138            val_loss, f1, acc = validate_epoch(model, val_loader, param)
139
140            # Save the best model and check for early stopping
141            best_fold_loss, patience_counter, overall_best_loss = save_best_model(
142                model,
143                val_loss,
144                best_fold_loss,
145                patience_counter,
146                overall_best_loss,
147                param,
148                fold,
149                best_model_path,
150            )
151            # Early stopping
152            epoch_pbar.set_postfix(
153                {
154                    "train_loss": f"{epoch_loss:.4f}",
155                    "val_loss": f"{val_loss:.4f}",
156                    "val_acc": f"{acc:.4f}",
157                    "val_f1": f"{f1:.4f}",
158                    "patience": patience_counter,
159                }
160            )
161
162            print(
163                f"Fold {fold+1} | Epoch {epoch+1} | Val Loss: {val_loss:.4f} | F1: {f1:.4f} | Acc: {acc:.4f}"
164            )
165
166            if patience_counter >= param["patience"]:
167                print(f"Early stopping at epoch {epoch+1} (fold {fold + 1})")
168                break
169
170    del model
171    empty_cache(param["device"])
172
173    print(f"\nTraining completed. Best model saved at: {best_model_path}")