zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Prefer MPS over CPU if available (#8210) 

* Prefer MPS over CPU if available

* Update torch_utils.py
This commit is contained in:
Glenn Jocher 2022-06-17 18:28:29 +02:00 committed by GitHub
parent 86e6add07b
commit 303f3785ae
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
hubconf.py
View File

@ -41,7 +41,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
name = Path(name) name = Path(name)
path = name.with_suffix('.pt') if name.suffix == '' and not name.is_dir() else name # checkpoint path path = name.with_suffix('.pt') if name.suffix == '' and not name.is_dir() else name # checkpoint path
try: try:
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device) device = select_device(device)
if pretrained and channels == 3 and classes == 80: if pretrained and channels == 3 and classes == 80:
model = DetectMultiBackend(path, device=device) # download/load FP32 model model = DetectMultiBackend(path, device=device) # download/load FP32 model

12
utils/torch_utils.py
View File

@ -62,8 +62,7 @@ def select_device(device='', batch_size=0, newline=True):
assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \ assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)" f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
cuda = not cpu and torch.cuda.is_available() if not cpu and torch.cuda.is_available(): # prefer GPU if available
if cuda:
devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7 devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7
n = len(devices) # device count n = len(devices) # device count
if n > 1 and batch_size > 0: # check batch_size is divisible by device_count if n > 1 and batch_size > 0: # check batch_size is divisible by device_count
@ -72,15 +71,18 @@ def select_device(device='', batch_size=0, newline=True):
for i, d in enumerate(devices): for i, d in enumerate(devices):
p = torch.cuda.get_device_properties(i) 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} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n" # bytes to MB
elif mps: arg = 'cuda:0'
elif not cpu and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available(): # prefer MPS if available
s += 'MPS\n' s += 'MPS\n'
else: arg = 'mps'
else: # revert to CPU
s += 'CPU\n' s += 'CPU\n'
arg = 'cpu'
if not newline: if not newline:
s = s.rstrip() s = s.rstrip()
LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
return torch.device('cuda:0' if cuda else 'mps' if mps else 'cpu') return torch.device(arg)
def time_sync(): def time_sync():