Mentors4EDU's picture
Publish training pipeline under training/; update repository notes
db3982a verified
Raw History Blame Contribute Delete
2.28 kB
from __future__ import annotations
import logging
import os
import torch
logger = logging.getLogger("riemann.device")
try:
import torch_directml
_DML_AVAILABLE = True
except ImportError:
torch_directml = None
_DML_AVAILABLE = False
def directml_available() -> bool:
return _DML_AVAILABLE and torch_directml.device_count() > 0
def main_device() -> torch.device:
if directml_available():
return torch_directml.device(0)
if torch.cuda.is_available():
return torch.device("cuda", 0)
logger.warning("No DirectML or CUDA device available; training will run on CPU.")
return torch.device("cpu")
def secondary_device() -> torch.device:
if directml_available() and torch_directml.device_count() > 1:
return torch_directml.device(1)
if torch.cuda.is_available() and torch.cuda.device_count() > 1:
return torch.device("cuda", 1)
return main_device()
def cpu_device() -> torch.device:
return torch.device("cpu")
def configure_cpu_threads(num_threads: int | None = None) -> int:
threads = num_threads or os.cpu_count() or 1
torch.set_num_threads(threads)
try:
torch.set_num_interop_threads(max(1, threads // 2))
except RuntimeError:
pass
return threads
def describe_devices() -> dict:
info = {
"directml_available": directml_available(),
"cuda_available": torch.cuda.is_available(),
"devices": [],
}
if directml_available():
for i in range(torch_directml.device_count()):
info["devices"].append(
{
"backend": "directml",
"index": i,
"name": torch_directml.device_name(i),
}
)
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
info["devices"].append(
{"backend": "cuda", "index": i, "name": torch.cuda.get_device_name(i)}
)
info["cpu_logical_cores"] = os.cpu_count()
info["cpu_threads_configured"] = torch.get_num_threads()
return info
if __name__ == "__main__":
import json
print(json.dumps(describe_devices(), indent=2))