Prefer MPS over CPU if available (#8210)
Browse files* Prefer MPS over CPU if available
* Update torch_utils.py
- hubconf.py +1 -1
- utils/torch_utils.py +7 -5
hubconf.py
CHANGED
@@ -41,7 +41,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
|
|
41 |
name = Path(name)
|
42 |
path = name.with_suffix('.pt') if name.suffix == '' and not name.is_dir() else name # checkpoint path
|
43 |
try:
|
44 |
-
device = select_device(
|
45 |
|
46 |
if pretrained and channels == 3 and classes == 80:
|
47 |
model = DetectMultiBackend(path, device=device) # download/load FP32 model
|
|
|
41 |
name = Path(name)
|
42 |
path = name.with_suffix('.pt') if name.suffix == '' and not name.is_dir() else name # checkpoint path
|
43 |
try:
|
44 |
+
device = select_device(device)
|
45 |
|
46 |
if pretrained and channels == 3 and classes == 80:
|
47 |
model = DetectMultiBackend(path, device=device) # download/load FP32 model
|
utils/torch_utils.py
CHANGED
@@ -62,8 +62,7 @@ def select_device(device='', batch_size=0, newline=True):
|
|
62 |
assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
|
63 |
f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
|
64 |
|
65 |
-
|
66 |
-
if cuda:
|
67 |
devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7
|
68 |
n = len(devices) # device count
|
69 |
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):
|
|
72 |
for i, d in enumerate(devices):
|
73 |
p = torch.cuda.get_device_properties(i)
|
74 |
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n" # bytes to MB
|
75 |
-
|
|
|
76 |
s += 'MPS\n'
|
77 |
-
|
|
|
78 |
s += 'CPU\n'
|
|
|
79 |
|
80 |
if not newline:
|
81 |
s = s.rstrip()
|
82 |
LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
|
83 |
-
return torch.device(
|
84 |
|
85 |
|
86 |
def time_sync():
|
|
|
62 |
assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
|
63 |
f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
|
64 |
|
65 |
+
if not cpu and torch.cuda.is_available(): # prefer GPU if available
|
|
|
66 |
devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7
|
67 |
n = len(devices) # device count
|
68 |
if n > 1 and batch_size > 0: # check batch_size is divisible by device_count
|
|
|
71 |
for i, d in enumerate(devices):
|
72 |
p = torch.cuda.get_device_properties(i)
|
73 |
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n" # bytes to MB
|
74 |
+
arg = 'cuda:0'
|
75 |
+
elif not cpu and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available(): # prefer MPS if available
|
76 |
s += 'MPS\n'
|
77 |
+
arg = 'mps'
|
78 |
+
else: # revert to CPU
|
79 |
s += 'CPU\n'
|
80 |
+
arg = 'cpu'
|
81 |
|
82 |
if not newline:
|
83 |
s = s.rstrip()
|
84 |
LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
|
85 |
+
return torch.device(arg)
|
86 |
|
87 |
|
88 |
def time_sync():
|