code/src/utils.py

"""Small training utilities: seeding, metrics, checkpointing, device."""
from __future__ import annotations

import os
import random

import numpy as np
import torch


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def resolve_device(pref: str = "auto") -> torch.device:
    if pref == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(pref)


class AverageMeter:
    """Tracks a running average (loss, accuracy, ...)."""

    def __init__(self) -> None:
        self.reset()

    def reset(self) -> None:
        self.sum = 0.0
        self.count = 0

    def update(self, value: float, n: int = 1) -> None:
        self.sum += float(value) * n
        self.count += n

    @property
    def avg(self) -> float:
        return self.sum / self.count if self.count else 0.0


@torch.no_grad()
def accuracy(logits: torch.Tensor, targets: torch.Tensor, topk=(1,)):
    """Return top-k accuracy (as fractions in [0, 1]) for each k."""
    maxk = max(topk)
    _, pred = logits.topk(maxk, dim=1, largest=True, sorted=True)  # (B, maxk)
    correct = pred.eq(targets.view(-1, 1))
    return [correct[:, :k].any(dim=1).float().mean().item() for k in topk]


def save_checkpoint(state: dict, out_dir: str, is_best: bool) -> None:
    os.makedirs(out_dir, exist_ok=True)
    last_path = os.path.join(out_dir, "last.pt")
    torch.save(state, last_path)
    if is_best:
        torch.save(state, os.path.join(out_dir, "best.pt"))