Spaces:
Sleeping
Sleeping
| from src.utils.logger import get_logger | |
| logger = get_logger(__name__) | |
| def get_device(): | |
| """ | |
| Detect best available device for PyTorch. | |
| Checks CUDA first, then MPS (Apple Silicon), then CPU. | |
| Returns: | |
| str: 'cuda', 'mps', or 'cpu' | |
| """ | |
| try: | |
| import torch | |
| if torch.cuda.is_available(): | |
| logger.info("CUDA available: %s", torch.cuda.get_device_name(0)) | |
| return "cuda" | |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| logger.info("MPS (Apple Silicon) available") | |
| return "mps" | |
| except ImportError: | |
| pass | |
| logger.info("Using CPU") | |
| return "cpu" | |
| def get_tf_device(): | |
| """ | |
| Detect best available device for TensorFlow. | |
| Returns: | |
| str: '/GPU:0' if GPU available, else '/CPU:0' | |
| """ | |
| try: | |
| import tensorflow as tf | |
| gpus = tf.config.list_physical_devices("GPU") | |
| if gpus: | |
| logger.info("TensorFlow GPUs: %s", gpus) | |
| return "/GPU:0" | |
| except ImportError: | |
| pass | |
| return "/CPU:0" | |
| def set_memory_growth(): | |
| """ | |
| Enable TensorFlow GPU memory growth to prevent OOM errors. | |
| Should be called before any TensorFlow operations. | |
| Safe to call even if no GPU is available. | |
| """ | |
| try: | |
| import tensorflow as tf | |
| gpus = tf.config.list_physical_devices("GPU") | |
| for gpu in gpus: | |
| tf.config.experimental.set_memory_growth(gpu, True) | |
| logger.info("Memory growth enabled for %s", gpu) | |
| except Exception as e: | |
| logger.warning("Could not set memory growth: %s", e) | |
| def log_device_info(): | |
| """ | |
| Log available hardware information for diagnostics. | |
| Logs PyTorch and TensorFlow device availability. | |
| """ | |
| try: | |
| import torch | |
| logger.info( | |
| "PyTorch version=%s CUDA=%s", | |
| torch.__version__, | |
| torch.cuda.is_available(), | |
| ) | |
| except ImportError: | |
| logger.info("PyTorch not available") | |
| try: | |
| import tensorflow as tf | |
| logger.info("TensorFlow version=%s", tf.__version__) | |
| gpus = tf.config.list_physical_devices("GPU") | |
| logger.info("TensorFlow GPUs: %d", len(gpus)) | |
| except ImportError: | |
| logger.info("TensorFlow not available") | |