From 9aca61c465e1162a1a5209059ff24b55c64760f4 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 13 Jul 2023 18:16:47 +0200 Subject: [PATCH] Print CPU information (#3721) --- ultralytics/yolo/utils/torch_utils.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/ultralytics/yolo/utils/torch_utils.py b/ultralytics/yolo/utils/torch_utils.py index 245f157..0ab1f4e 100644 --- a/ultralytics/yolo/utils/torch_utils.py +++ b/ultralytics/yolo/utils/torch_utils.py @@ -18,7 +18,7 @@ import torch.nn.functional as F import torchvision from ultralytics.yolo.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, __version__ -from ultralytics.yolo.utils.checks import check_version +from ultralytics.yolo.utils.checks import check_requirements, check_version try: import thop @@ -53,6 +53,13 @@ def smart_inference_mode(): return decorate +def get_cpu_info(): + """Return a string with system CPU information, i.e. 'Apple M2'.""" + check_requirements('py-cpuinfo') + import cpuinfo # noqa + return cpuinfo.get_cpu_info()['brand_raw'].replace('(R)', '').replace('CPU ', '').replace('@ ', '') + + def select_device(device='', batch=0, newline=False, verbose=True): """Selects PyTorch Device. Options are device = None or 'cpu' or 0 or '0' or '0,1,2,3'.""" s = f'Ultralytics YOLOv{__version__} 🚀 Python-{platform.python_version()} torch-{torch.__version__} ' @@ -93,10 +100,10 @@ def select_device(device='', batch=0, newline=False, verbose=True): arg = 'cuda:0' elif mps and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available() and TORCH_2_0: # Prefer MPS if available - s += 'MPS\n' + s += f'MPS ({get_cpu_info()})\n' arg = 'mps' else: # revert to CPU - s += 'CPU\n' + s += f'CPU ({get_cpu_info()})\n' arg = 'cpu' if verbose and RANK == -1: