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