mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
"""Reproducibility utilities for deterministic abliteration runs.
|
|
|
|
Sets all random seeds (Python, NumPy, PyTorch CPU/CUDA) and optionally
|
|
enables PyTorch deterministic mode for bit-exact reproducibility.
|
|
|
|
Usage:
|
|
from obliteratus.reproducibility import set_seed
|
|
|
|
set_seed(42) # sets all seeds, enables deterministic ops
|
|
set_seed(42, deterministic=False) # seeds only, faster
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import random
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def set_seed(seed: int = 42, deterministic: bool = True) -> None:
|
|
"""Set all random seeds for reproducibility.
|
|
|
|
Args:
|
|
seed: The seed value to use everywhere.
|
|
deterministic: If True, also enable PyTorch deterministic algorithms
|
|
and disable cuDNN benchmarking. This is slower but guarantees
|
|
bit-exact results across runs on the same hardware.
|
|
"""
|
|
random.seed(seed)
|
|
|
|
try:
|
|
import numpy as np
|
|
np.random.seed(seed)
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
import torch
|
|
from obliteratus import device as dev
|
|
torch.manual_seed(seed)
|
|
dev.set_seed_all(seed)
|
|
|
|
if deterministic:
|
|
torch.use_deterministic_algorithms(True, warn_only=True)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
|
except ImportError:
|
|
pass
|
|
|
|
logger.debug("Seeds set to %d (deterministic=%s)", seed, deterministic)
|