glenn-jocher commited on
Commit
303f378
·
unverified ·
1 Parent(s): 86e6add

Prefer MPS over CPU if available (#8210)

Browse files

* Prefer MPS over CPU if available

* Update torch_utils.py

Files changed (2) hide show
  1. hubconf.py +1 -1
  2. 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(('0' if torch.cuda.is_available() else 'cpu') if device is None else 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
- cuda = not cpu and torch.cuda.is_available()
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
- elif mps:
 
76
  s += 'MPS\n'
77
- else:
 
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('cuda:0' if cuda else 'mps' if mps else 'cpu')
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():