File size: 1,787 Bytes
05b535a | 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 | """Utility helper functions."""
import random
import torch
import numpy as np
def set_seed(seed: int) -> None:
"""Set random seed for reproducibility."""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def get_device(device_str: str = "cuda") -> torch.device:
"""
Get torch device.
Args:
device_str: Device string ("cuda" or "cpu")
Returns:
torch.device instance
"""
if device_str == "cuda" and torch.cuda.is_available():
return torch.device("cuda")
return torch.device("cpu")
def format_size(num_bytes: int) -> str:
"""Format bytes to human-readable size."""
for unit in ["B", "KB", "MB", "GB"]:
if num_bytes < 1024:
return f"{num_bytes:.1f}{unit}"
num_bytes /= 1024
return f"{num_bytes:.1f}TB"
def count_parameters(model: torch.nn.Module) -> tuple[int, int]:
"""
Count parameters in model.
Args:
model: PyTorch model
Returns:
Tuple of (total_params, trainable_params)
"""
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
return total_params, trainable_params
def get_dtype(dtype_str: str) -> torch.dtype:
"""
Get torch dtype from string.
Args:
dtype_str: Dtype string (float32, float16, bfloat16)
Returns:
torch.dtype
"""
dtype_map = {
"float32": torch.float32,
"float16": torch.float16,
"bfloat16": torch.bfloat16,
}
return dtype_map.get(dtype_str, torch.float32)
|