|
|
|
@ -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: |
|
|
|
|