Droid210/FleetVision

1"""Utility functions for device management and seed setting."""2import random3 4import numpy as np5import torch6 7 8def set_seed(seed: int) -> None:9    """Set random seed for reproducibility.10 11    Args:12        seed: Random seed value.13    """14    random.seed(seed)15    np.random.seed(seed)16    torch.manual_seed(seed)17    torch.cuda.manual_seed_all(seed)18 19 20def get_device() -> torch.device:21    """Get available device with fallback to CPU.22 23    Returns:24        torch.device: CUDA if available, otherwise CPU.25    """26    if torch.cuda.is_available():27        try:28            # Test CUDA availability29            torch.cuda.get_device_name(0)30            # Test with Conv2d like the ViT model will use31            from torch import nn32            test_input = torch.randn(1, 3, 32, 32).cuda()33            test_conv = nn.Conv2d(3, 16, kernel_size=3).cuda()34            _ = test_conv(test_input)35            del test_input, test_conv36            torch.cuda.empty_cache()37            return torch.device("cuda")38        except RuntimeError:39            print("WARNING: CUDA detected but not usable. Using CPU.")40            return torch.device("cpu")41    return torch.device("cpu")42