HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.3.6 improve collect_system_info (#16722)

Browse source 
Signed-off-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-10-06 20:36:35 +02:00 • committed by GitHub
parent 4884311991
commit d90efbd297
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 52 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

9
ultralytics/utils/torch_utils.py
View file

@ -123,6 +123,12 @@ def get_cpu_info():
return PERSISTENT_CACHE.get("cpu_info", "unknown")
def get_gpu_info(index):
"""Return a string with system GPU information, i.e. 'Tesla T4, 15102MiB'."""
properties = torch.cuda.get_device_properties(index)
return f"{properties.name}, {properties.total_memory / (1 << 20):.0f}MiB"
def select_device(device="", batch=0, newline=False, verbose=True):
"""
Selects the appropriate PyTorch device based on the provided arguments.
@ -208,8 +214,7 @@ def select_device(device="", batch=0, newline=False, verbose=True):
)
space = " " * (len(s) + 1)
for i, d in enumerate(devices):
p = torch.cuda.get_device_properties(i)
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n" # bytes to MB
s += f"{'' if i == 0 else space}CUDA:{d} ({get_gpu_info(i)})\n" # bytes to MB
arg = "cuda:0"
elif mps and TORCH_2_0 and torch.backends.mps.is_available():
# Prefer MPS if available