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