File size: 943 Bytes
9466fff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
"""
src/utils.py

Small shared helpers. Nothing domain-specific lives here — just things that
would otherwise be copy-pasted across train.py / evaluate.py.
"""

import random

import numpy as np
import torch


def set_seed(seed: int) -> None:
    """Make training reproducible across runs with the same seed."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    # Slight performance hit but needed for strict reproducibility
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def get_device() -> torch.device:
    """Return CUDA if available, else CPU."""
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_yaml(path: str) -> dict:
    """Load a YAML config file into a plain dict."""
    import yaml
    with open(path, "r") as f:
        return yaml.safe_load(f)