| |
| |
| |
| |
| |
| |
| |
| import torch |
| import logging |
|
|
|
|
| def get_device(number: int, logger: logging.Logger = None): |
| """ |
| Configures PyTorch to use a specified GPU by its index number, |
| or falls back to CPU if CUDA is not available. |
| |
| Args: |
| number (int): The index number of the GPU to use. |
| logger (logging.Logger, optional): Logger for logging GPU info. |
| |
| Returns: |
| torch.device: The selected torch device (GPU or CPU). |
| """ |
| |
| if not torch.cuda.is_available(): |
| if logger: |
| logger.warning("CUDA is not available. Falling back to CPU.") |
| return torch.device('cpu') |
|
|
| |
| if number >= torch.cuda.device_count() or number < 0: |
| raise ValueError( |
| f"GPU number {number} is not valid. Available GPU indices range from 0 to {torch.cuda.device_count() - 1}.") |
|
|
| |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats() |
| torch.cuda.reset_accumulated_memory_stats() |
|
|
| |
| torch.cuda.set_device(number) |
| if logger: |
| logger.info(f"PyTorch is now configured to use GPU {number}: {torch.cuda.get_device_name(number)}") |
|
|
| device_name = torch.cuda.get_device_name(number) |
| total_mem = torch.cuda.get_device_properties(number).total_memory / 1024 ** 2 |
| mem_allocated = torch.cuda.memory_allocated(number) / 1024 ** 2 |
| mem_reserved = torch.cuda.memory_reserved(number) / 1024 ** 2 |
| max_allocated = torch.cuda.max_memory_allocated(number) / 1024 ** 2 |
| max_reserved = torch.cuda.max_memory_reserved(number) / 1024 ** 2 |
|
|
| logger.info(f"[GPU {number} - {device_name}] Memory Stats:") |
| logger.info(f" Total Memory : {total_mem:.2f} MB") |
| logger.info(f" Currently Allocated : {mem_allocated:.2f} MB") |
| logger.info(f" Currently Reserved : {mem_reserved:.2f} MB") |
| logger.info(f" Max Allocated : {max_allocated:.2f} MB") |
| logger.info(f" Max Reserved : {max_reserved:.2f} MB") |
|
|
| return torch.device(f'cuda:{number}') |
| |
| |
| |
|
|
