garmentiq.classification.test_pytorch_nn

  1"""Evaluating a classification model on a held-out dataset."""
  2import torch
  3import torch.nn as nn
  4from torch.utils.data import DataLoader, TensorDataset
  5from typing import Callable, Type
  6from tqdm.auto import tqdm
  7from sklearn.metrics import f1_score, accuracy_score, classification_report
  8from garmentiq.utils.device import empty_cache
  9from garmentiq.classification.utils import (
 10    CachedDataset,
 11    seed_worker,
 12    train_epoch,
 13    validate_epoch,
 14    save_best_model,
 15    validate_train_param,
 16    validate_test_param,
 17)
 18
 19
 20def test_pytorch_nn(
 21    model_path: str,
 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    Evaluates a trained PyTorch model on a test dataset.
 30
 31    Loads the model from disk, prepares the test dataset, and computes loss, accuracy,
 32    F1 score, and prints a full classification report.
 33
 34    Args:
 35        model_path (str): Path to the saved model checkpoint file.
 36        model_class (Type[torch.nn.Module]): The class of the PyTorch model to instantiate.
 37                                            Must inherit from `torch.nn.Module`.
 38        model_args (dict): Dictionary of arguments used to initialize the model.
 39        dataset_class (Callable): A callable class or function that returns a `torch.utils.data.Dataset`-compatible dataset.
 40                                  (Note: Not directly used, but included for consistency with training pipeline.)
 41        dataset_args (dict): Dictionary with dataset components:
 42            - `cached_images` (torch.Tensor): Preprocessed test image tensors.
 43            - `cached_labels` (torch.Tensor): Corresponding test labels.
 44            - `raw_labels` (pandas.Series or array-like): Original labels for report generation.
 45        param (dict): Dictionary of optional configuration parameters.
 46                      Optional Keys:
 47                          - `device` (Union[str, torch.device]): Device for computation, e.g. `"cpu"`,
 48                            `"cuda"`, `"cuda:0"`, or `"mps"`. Hardware acceleration is opt-in;
 49                            pass it explicitly to use a GPU or Apple Silicon. Default is `"cpu"`.
 50                          - `batch_size` (int): Batch size used for testing. Default is 64.
 51
 52    Raises:
 53        FileNotFoundError: If the model checkpoint cannot be loaded.
 54        ValueError: If the requested `device` is invalid or unavailable on this machine.
 55        TypeError: If any parameter is of an incorrect type.
 56
 57    Returns:
 58        None — prints test loss, accuracy, F1 score, and a classification report.
 59    """
 60    validate_test_param(param)
 61    model = model_class(**model_args).to(param["device"])
 62    state_dict = torch.load(model_path, map_location=param["device"], weights_only=True)
 63    new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
 64    model.load_state_dict(new_state_dict, strict=False)
 65    model.eval()
 66
 67    test_dataset = TensorDataset(
 68        dataset_args["cached_images"], dataset_args["cached_labels"]
 69    )
 70    test_loader = DataLoader(
 71        test_dataset, batch_size=param["batch_size"], shuffle=False
 72    )
 73
 74    # Evaluation
 75    all_preds = []
 76    all_labels = []
 77    total_loss = 0.0
 78
 79    with torch.no_grad():
 80        for images, labels in tqdm(test_loader, desc="Evaluating"):
 81            images = images.to(param["device"])
 82            labels = labels.to(param["device"])
 83
 84            outputs = model(images)
 85            criterion = nn.CrossEntropyLoss()
 86            loss = criterion(outputs, labels)
 87
 88            total_loss += loss.item() * images.size(0)
 89            _, preds = torch.max(outputs, 1)
 90
 91            all_preds.extend(preds.cpu().numpy())
 92            all_labels.extend(labels.cpu().numpy())
 93
 94    # Calculate metrics
 95    test_loss = total_loss / len(test_loader.dataset)
 96    test_acc = accuracy_score(all_labels, all_preds)
 97    test_f1 = f1_score(all_labels, all_preds, average="weighted")
 98
 99    del model
100    empty_cache(param["device"])
101
102    print(f"Test Loss: {test_loss:.4f}")
103    print(f"Test Accuracy: {test_acc:.4f}")
104    print(f"Test F1 Score: {test_f1:.4f}")
105    print("\nClassification Report:")
106    print(
107        classification_report(
108            all_labels,
109            all_preds,
110            target_names=sorted(dataset_args["raw_labels"].unique()),
111        )
112    )