Browse Source

Multi-GPU default to single device 0 (#3554)

* Multi-GPU default to single device 0

* Multi-GPU default to single device 0

* add space
modifyDataloader
Glenn Jocher GitHub 3 years ago
parent
commit
8b5086c21b
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
  1. +2
    -2
      utils/torch_utils.py

+ 2
- 2
utils/torch_utils.py View File



cuda = not cpu and torch.cuda.is_available() cuda = not cpu and torch.cuda.is_available()
if cuda: if cuda:
devices = device.split(',') if device else 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: # check batch_size is divisible by device_count if n > 1 and batch_size: # check batch_size is divisible by device_count
assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}' assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
space = ' ' * len(s)
space = ' ' * (len(s) + 1)
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 / 1024 ** 2}MB)\n" # bytes to MB s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB

Loading…
Cancel
Save