Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-09-26 13:52:37 +02:00
parent f146692ad0
commit df8529a747
1 changed files with 10 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
utils/torch_utils.py
View File

@ -15,33 +15,28 @@ def init_seeds(seed=0):
def select_device(device='', apex=False): def select_device(device='', apex=False):
# device = '' or 'cpu' or '0' or '0,1' # device = 'cpu' or '0' or '0,1,2,3'
if device.lower() == 'cpu': cpu_request = device.lower() == 'cpu'
pass if device and not cpu_request: # if device requested other than 'cpu'
elif device: # Set environment variable if device is specified os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity
os.environ['CUDA_VISIBLE_DEVICES'] = device
# apex if mixed precision training https://github.com/NVIDIA/apex cuda = False if cpu_request else torch.cuda.is_available()
cuda = False if device == 'cpu' else torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')
if not cuda:
print('Using CPU')
if cuda: if cuda:
c = 1024 ** 2 # bytes to MB c = 1024 ** 2 # bytes to MB
ng = torch.cuda.device_count() ng = torch.cuda.device_count()
x = [torch.cuda.get_device_properties(i) for i in range(ng)] x = [torch.cuda.get_device_properties(i) for i in range(ng)]
cuda_str = 'Using CUDA ' + ('Apex ' if apex else '') cuda_str = 'Using CUDA ' + ('Apex ' if apex else '') # apex for mixed precision https://github.com/NVIDIA/apex
for i in range(0, ng): for i in range(0, ng):
if i == 1: if i == 1:
# torch.cuda.set_device(0) # OPTIONAL: Set GPU ID
cuda_str = ' ' * len(cuda_str) cuda_str = ' ' * len(cuda_str)
print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
(cuda_str, i, x[i].name, x[i].total_memory / c)) (cuda_str, i, x[i].name, x[i].total_memory / c))
else:
print('Using CPU')
print('') # skip a line print('') # skip a line
return device return torch.device('cuda:0' if cuda else 'cpu')
def fuse_conv_and_bn(conv, bn): def fuse_conv_and_bn(conv, bn):