alnn_rewrite/util/seeding.py

"""Centralized seeding helpers for reproducible training and evaluation.

Reproducibility matters here because the whole point of the harness is to
compare ensemble members and quantify uncertainty. If weight initialization,
dropout, dataloader shuffling, and BNN sampling are not seeded, "ensemble
diversity" is accidental and results are not reproducible across runs.
"""

import os
import random
from typing import Optional

import numpy as np
import torch

# Numpy requires seeds in the uint32 range; torch/python are more permissive,
# so we clamp everything to this range for consistency.
_UINT32 = 2**32


def seed_everything(seed: int, deterministic: bool = False) -> int:
    """Seed all relevant RNGs (python, numpy, torch, CUDA).

    Args:
        seed: The master seed. Clamped into the uint32 range.
        deterministic: If True, force cuDNN into deterministic mode. This makes
            convolutions reproducible at some throughput cost. Leave False for
            training speed; set True when exact reproduction is required.

    Returns:
        The (clamped) seed actually applied, for logging.
    """
    seed = int(seed) % _UINT32

    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        # Opt into deterministic algorithms where available; warn_only so a
        # missing deterministic kernel degrades gracefully instead of crashing.
        torch.use_deterministic_algorithms(True, warn_only=True)

    return seed


def derive_seed(base_seed: int, index: int, stream: int = 0) -> int:
    """Derive a stable, distinct seed for ensemble member ``index``.

    ``stream`` separates otherwise-colliding families (e.g. normal vs. Bayesian
    ensembles) so that normal-member-0 and bayesian-member-0 do not share an
    identical RNG stream.

    Args:
        base_seed: The experiment master seed.
        index: The ensemble member index (0-based).
        stream: An offset identifying the RNG family.

    Returns:
        A seed in the uint32 range, deterministic in all three arguments.
    """
    return (int(base_seed) + stream * 100_000 + int(index)) % _UINT32


def torch_generator(seed: Optional[int]) -> Optional[torch.Generator]:
    """Return a seeded ``torch.Generator`` for DataLoader shuffling, or None.

    Giving the training DataLoader its own generator makes the shuffle order a
    pure function of ``seed`` rather than of global RNG state mutated elsewhere.
    """
    if seed is None:
        return None
    return torch.Generator().manual_seed(int(seed) % _UINT32)