File size: 2,978 Bytes
d7228c8 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """
General utility functions for seed setting, device resolution, and parameter formatting.
"""
import os
import random
from typing import Dict, Any, Union
import numpy as np
import torch
def set_seed(seed: int = 42) -> None:
"""
Sets reproducible random seeds across Python random, NumPy, PyTorch CPU, and PyTorch CUDA.
Args:
seed: Integer seed value.
"""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
os.environ["PYTHONHASHSEED"] = str(seed)
def get_device(device_arg: str = "auto") -> torch.device:
"""
Resolves execution device based on hardware availability.
Args:
device_arg: Device choice ('auto', 'cuda', 'cpu', 'mps').
Returns:
torch.device instance.
"""
if device_arg == "auto":
if torch.cuda.is_available():
return torch.device("cuda")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps")
else:
return torch.device("cpu")
return torch.device(device_arg)
def count_parameters(model: torch.nn.Module) -> Dict[str, int]:
"""
Calculates total and trainable parameter count of a PyTorch module.
Args:
model: Target PyTorch module.
Returns:
Dictionary containing total, trainable, and non-trainable parameter counts.
"""
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
return {
"total": total,
"trainable": trainable,
"non_trainable": total - trainable
}
def format_number(num: Union[int, float]) -> str:
"""
Formats large numbers into human-readable metric prefixes (K, M, B).
Args:
num: Input numerical value.
Returns:
Formatted string (e.g. 1.23M).
"""
if num >= 1e9:
return f"{num / 1e9:.2f}B"
if num >= 1e6:
return f"{num / 1e6:.2f}M"
if num >= 1e3:
return f"{num / 1e3:.2f}K"
return str(num)
def get_memory_stats(device: torch.device) -> Dict[str, Any]:
"""
Retrieves CUDA GPU memory usage statistics if available.
Args:
device: Active torch device.
Returns:
Dictionary of allocated and reserved VRAM in MB.
"""
if device.type == "cuda" and torch.cuda.is_available():
allocated = torch.cuda.memory_allocated(device) / (1024 ** 2)
reserved = torch.cuda.memory_reserved(device) / (1024 ** 2)
max_allocated = torch.cuda.max_memory_allocated(device) / (1024 ** 2)
return {
"allocated_mb": round(allocated, 2),
"reserved_mb": round(reserved, 2),
"max_allocated_mb": round(max_allocated, 2)
}
return {"allocated_mb": 0.0, "reserved_mb": 0.0, "max_allocated_mb": 0.0}
|