garmentiq.classification.utils

  1"""Training helpers for the classification module.
  2
  3Holds the in-memory dataset wrapper, the per-epoch train and validation loops,
  4checkpoint saving, and the parameter validators that apply defaults to the `param`
  5dictionary, including normalising `device`.
  6"""
  7import torch
  8import torch.nn as nn
  9import torch.optim as optim
 10from torch.utils.data import DataLoader, Dataset
 11from typing import Callable
 12from tqdm.auto import tqdm
 13import os
 14from sklearn.metrics import f1_score, accuracy_score
 15import random
 16import numpy as np
 17from garmentiq.utils.device import resolve_device
 18
 19
 20class CachedDataset(Dataset):
 21    """
 22    A PyTorch Dataset that wraps pre-loaded data (images and labels) in memory.
 23
 24    This dataset is designed to be used when images and labels have already been
 25    loaded and preprocessed into PyTorch tensors or NumPy arrays, avoiding
 26    repeated disk I/O during training/validation.
 27
 28    Attributes:
 29        indices (list or numpy.ndarray): A list or array of indices that map
 30            to specific items in `cached_images` and `cached_labels`. This
 31            allows for flexible subsetting (e.g., for train/validation splits).
 32        cached_images (torch.Tensor): A tensor containing the pre-loaded images.
 33        cached_labels (torch.Tensor): A tensor containing the pre-loaded labels.
 34    """
 35
 36    def __init__(self, indices, cached_images, cached_labels):
 37        """
 38        Initializes the CachedDataset.
 39
 40        Args:
 41            indices (list or numpy.ndarray): Indices to select from the cached data.
 42            cached_images (torch.Tensor): Pre-loaded image tensor.
 43            cached_labels (torch.Tensor): Pre-loaded label tensor.
 44        """
 45        self.indices = indices
 46        self.cached_images = cached_images
 47        self.cached_labels = cached_labels
 48
 49    def __len__(self):
 50        """
 51        Returns the number of samples in the dataset.
 52
 53        Returns:
 54            int: The number of samples.
 55        """
 56        return len(self.indices)
 57
 58    def __getitem__(self, idx):
 59        """
 60        Retrieves a sample from the dataset at the given index.
 61
 62        Args:
 63            idx (int): The index of the sample to retrieve.
 64
 65        Returns:
 66            tuple: A tuple containing the image and its corresponding label.
 67        """
 68        actual_idx = self.indices[idx]
 69        return self.cached_images[actual_idx], self.cached_labels[actual_idx]
 70
 71
 72def seed_worker(worker_id, SEED=88):
 73    """
 74    Seeds the random number generators for a DataLoader worker.
 75
 76    This function is intended to be passed as `worker_init_fn` to a PyTorch
 77    DataLoader to ensure reproducibility across different worker processes.
 78    It seeds Python's `random` module, NumPy, and PyTorch for each worker.
 79
 80    Args:
 81        worker_id (int): The ID of the current worker process.
 82        SEED (int, optional): The base seed value. The worker's seed will be
 83            `SEED + worker_id`. Defaults to 88.
 84    """
 85    worker_seed = SEED + worker_id
 86    random.seed(worker_seed)
 87    np.random.seed(worker_seed)
 88    torch.manual_seed(worker_seed)
 89
 90
 91def train_epoch(model, train_loader, optimizer, param):
 92    """
 93    Performs a single training epoch for a PyTorch model.
 94
 95    Sets the model to training mode, iterates through the `train_loader`,
 96    performs forward and backward passes, and updates model weights using
 97    the provided optimizer. A progress bar is displayed, showing the batch loss.
 98
 99    Args:
100        model (torch.nn.Module): The PyTorch model to train.
101        train_loader (torch.utils.data.DataLoader): DataLoader for the training data.
102        optimizer (torch.optim.Optimizer): The optimizer used for updating model weights.
103        param (dict): A dictionary containing training parameters, including:
104            - "device" (torch.device): The device to use for training, e.g. `"cpu"`,
105              `"cuda"`, `"cuda:0"`, or `"mps"`. Normalised by `validate_train_param`.
106
107    Returns:
108        float: The average loss for the epoch.
109    """
110    model.train()
111    running_loss = 0.0
112    train_pbar = tqdm(train_loader, desc=f"Training", leave=False)
113
114    for images, labels in train_pbar:
115        images, labels = images.to(param["device"], non_blocking=True), labels.to(
116            param["device"], non_blocking=True
117        )
118        optimizer.zero_grad()
119        outputs = model(images)
120
121        # Calculate the loss
122        criterion = nn.CrossEntropyLoss()
123        loss = criterion(outputs, labels)
124        loss.backward()
125        optimizer.step()
126
127        running_loss += loss.item() * images.size(0)
128        train_pbar.set_postfix({"batch_loss": f"{loss.item():.4f}"})
129
130    epoch_loss = running_loss / len(train_loader.dataset)
131    return epoch_loss
132
133
134def validate_epoch(model, val_loader, param):
135    """
136    Performs a single validation epoch for a PyTorch model.
137
138    Sets the model to evaluation mode, iterates through the `val_loader`
139    without gradient calculations, and computes the validation loss, F1 score,
140    and accuracy. A progress bar is displayed, showing the batch validation loss.
141
142    Args:
143        model (torch.nn.Module): The PyTorch model to validate.
144        val_loader (torch.utils.data.DataLoader): DataLoader for the validation data.
145        param (dict): A dictionary containing training parameters, including:
146            - "device" (torch.device): The device to use for validation, e.g. `"cpu"`,
147              `"cuda"`, `"cuda:0"`, or `"mps"`. Normalised by `validate_train_param`.
148
149    Returns:
150        tuple: A tuple containing:
151            - val_loss (float): The average validation loss for the epoch.
152            - f1 (float): The weighted average F1 score on the validation set.
153            - acc (float): The accuracy score on the validation set.
154    """
155    model.eval()
156    val_loss = 0.0
157    all_preds = []
158    all_labels = []
159    val_pbar = tqdm(val_loader, desc=f"Validation", leave=False)
160
161    with torch.no_grad():
162        for images, labels in val_pbar:
163            images, labels = images.to(param["device"]), labels.to(param["device"])
164            outputs = model(images)
165            criterion = nn.CrossEntropyLoss()
166            loss = criterion(outputs, labels)
167            val_loss += loss.item() * images.size(0)
168
169            # Collect predictions and true labels for F1 score calculation
170            _, preds = torch.max(outputs, 1)
171            all_preds.extend(preds.cpu().numpy())  # Move to CPU and convert to numpy
172            all_labels.extend(labels.cpu().numpy())  # Move to CPU and convert to numpy
173
174            val_pbar.set_postfix({"val_batch_loss": f"{loss.item():.4f}"})
175
176    # Calculate the average validation loss
177    val_loss /= len(val_loader.dataset)
178
179    # Calculate F1 Score (if needed)
180    f1 = f1_score(all_labels, all_preds, average="weighted")
181    acc = accuracy_score(all_labels, all_preds)
182
183    return val_loss, f1, acc
184
185
186def save_best_model(
187    model,
188    val_loss,
189    best_fold_loss,
190    patience_counter,
191    overall_best_loss,
192    param,
193    fold,
194    best_model_path,
195):
196    """
197    Saves the best model checkpoints based on validation loss and manages early stopping.
198
199    This function updates the `best_fold_loss` and `patience_counter` for the current
200    cross-validation fold. It also saves the model's state dictionary if it's the best
201    performing model for the current fold or the overall best model across all folds.
202
203    Args:
204        model (torch.nn.Module): The current PyTorch model being trained.
205        val_loss (float): The validation loss from the current epoch.
206        best_fold_loss (float): The best validation loss recorded so far for the current fold.
207        patience_counter (int): The number of epochs since the last improvement for the current fold.
208        overall_best_loss (float): The best validation loss recorded so far across all folds.
209        param (dict): A dictionary containing training parameters, including:
210            - "model_save_dir" (str): Directory where model checkpoints will be saved.
211        fold (int): The current fold number (0-indexed).
212        best_model_path (str): The full path where the overall best model will be saved.
213
214    Returns:
215        tuple: A tuple containing:
216            - best_fold_loss (float): The updated best validation loss for the current fold.
217            - patience_counter (int): The updated patience counter for the current fold.
218            - overall_best_loss (float): The updated overall best validation loss.
219    """
220    # Save the best model for this fold
221    if val_loss < best_fold_loss:
222        best_fold_loss = val_loss
223        patience_counter = 0
224        torch.save(
225            model.state_dict(),
226            os.path.join(param["model_save_dir"], f"fold_{fold + 1}_best.pt"),
227        )
228    else:
229        patience_counter += 1
230
231    # Save the overall best model
232    if val_loss < overall_best_loss:
233        overall_best_loss = val_loss
234        torch.save(model.state_dict(), best_model_path)
235
236    return best_fold_loss, patience_counter, overall_best_loss
237
238
239def validate_train_param(param: dict):
240    """
241    Validates the parameter dictionary for training configuration.
242
243    This function checks for the presence and correct types of required
244    parameters for training, and applies default values for optional parameters
245    if they are not provided. The `device` entry is normalised in place into a
246    `torch.device`, so callers may supply either a string or a `torch.device`.
247
248    Args:
249        param (dict): The dictionary of training parameters to validate.
250                      Optional Keys:
251                          - `device` (Union[str, torch.device]): The device to train on, e.g.
252                            `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware acceleration is
253                            opt-in; pass it explicitly to use a GPU or Apple Silicon.
254                            Default is `"cpu"`.
255
256    Raises:
257        ValueError: If a required parameter is missing, or if the requested `device` is invalid
258                    or unavailable on this machine.
259        TypeError: If a parameter has an incorrect type.
260    """
261    # --- Required fields and types
262    required_keys = {"optimizer_class": type, "optimizer_args": dict}
263
264    # --- Normalise the requested device (accepts str or torch.device, defaults to CPU)
265    param["device"] = resolve_device(param.get("device", "cpu"))
266
267    # --- Optional fields with default values and expected types
268    optional_keys = {
269        "n_fold": (int, 5),
270        "n_epoch": (int, 100),
271        "patience": (int, 5),
272        "batch_size": (int, 64),
273        "model_save_dir": (str, "./models"),
274        "seed": (int, 88),
275        "seed_worker": (Callable, seed_worker),
276        "max_workers": (int, 0),
277        "best_model_name": (str, "best_model.pt"),
278        "pin_memory": (bool, False),
279        "persistent_workers": (bool, False),
280    }
281
282    # --- Validate required keys
283    for key, expected_types in required_keys.items():
284        if key not in param:
285            raise ValueError(f"Missing required param key: '{key}'")
286        if not isinstance(param[key], expected_types):
287            raise TypeError(
288                f"param['{key}'] must be of type {expected_types}, got {type(param[key])}"
289            )
290
291    # --- Apply defaults and type-check optional keys
292    for key, (expected_type, default) in optional_keys.items():
293        if key not in param:
294            param[key] = default
295        elif expected_type is not None and not isinstance(param[key], expected_type):
296            raise TypeError(
297                f"param['{key}'] must be of type {expected_type}, got {type(param[key])}"
298            )
299
300
301def validate_test_param(param: dict):
302    """
303    Validates the parameter dictionary for testing configuration.
304
305    This function checks for the presence and correct types of optional
306    parameters for testing, and applies default values if they are not provided.
307    The `device` entry is normalised in place into a `torch.device`, so callers
308    may supply either a string or a `torch.device`.
309
310    Args:
311        param (dict): The dictionary of testing parameters to validate.
312                      Optional Keys:
313                          - `device` (Union[str, torch.device]): The device to evaluate on, e.g.
314                            `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware acceleration is
315                            opt-in; pass it explicitly to use a GPU or Apple Silicon.
316                            Default is `"cpu"`.
317
318    Raises:
319        ValueError: If the requested `device` is invalid or unavailable on this machine.
320        TypeError: If a parameter has an incorrect type.
321    """
322    # --- Normalise the requested device (accepts str or torch.device, defaults to CPU)
323    param["device"] = resolve_device(param.get("device", "cpu"))
324
325    # --- Optional fields with default values and expected types
326    optional_keys = {
327        "batch_size": (int, 64),
328    }
329
330    # --- Apply defaults and type-check optional keys
331    for key, (expected_type, default) in optional_keys.items():
332        if key not in param:
333            param[key] = default
334        elif expected_type is not None and not isinstance(param[key], expected_type):
335            raise TypeError(
336                f"param['{key}'] must be of type {expected_type}, got {type(param[key])}"
337            )
338
339
340def validate_pred_param(param: dict):
341    """
342    Validates the parameter dictionary for prediction configuration.
343
344    This function checks for the presence and correct types of optional
345    parameters for prediction, and applies default values if they are not provided.
346    The `device` entry is normalised in place into a `torch.device`, so callers
347    may supply either a string or a `torch.device`.
348
349    Args:
350        param (dict): The dictionary of prediction parameters to validate.
351                      Optional Keys:
352                          - `device` (Union[str, torch.device]): The device to predict on, e.g.
353                            `"cpu"`, `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware acceleration is
354                            opt-in; pass it explicitly to use a GPU or Apple Silicon.
355                            Default is `"cpu"`.
356
357    Raises:
358        ValueError: If the requested `device` is invalid or unavailable on this machine.
359        TypeError: If a parameter has an incorrect type.
360    """
361    # --- Normalise the requested device (accepts str or torch.device, defaults to CPU)
362    param["device"] = resolve_device(param.get("device", "cpu"))
363
364    # --- Optional fields with default values and expected types
365    optional_keys = {
366        "batch_size": (int, 64),
367    }
368
369    # --- Apply defaults and type-check optional keys
370    for key, (expected_type, default) in optional_keys.items():
371        if key not in param:
372            param[key] = default
373        elif expected_type is not None and not isinstance(param[key], expected_type):
374            raise TypeError(
375                f"param['{key}'] must be of type {expected_type}, got {type(param[key])}"
376            )