Spaces:
Sleeping
Sleeping
File size: 2,308 Bytes
e382dce | 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 | 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")
|