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

updates

This commit is contained in:
Glenn Jocher 2019-07-16 18:58:49 +02:00
parent 64b606a3cd
commit 153762dec0
2 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
train.py
View File

@ -156,7 +156,6 @@ def train(cfg,
try: try:
from apex import amp from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0) model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)
print('Using Apex')
except: # not installed: install help: https://github.com/NVIDIA/apex/issues/259 except: # not installed: install help: https://github.com/NVIDIA/apex/issues/259
mixed_precision = False mixed_precision = False

15
utils/torch_utils.py
View File

@ -16,17 +16,24 @@ def select_device(force_cpu=False):
if not cuda: if not cuda:
print('Using CPU') print('Using CPU')
if cuda: if cuda:
try: # Mixed precision training https://github.com/NVIDIA/apex
from apex import amp
apex_str = ' with Apex'
except:
apex_str = ''
torch.backends.cudnn.benchmark = True # set False for reproducible results torch.backends.cudnn.benchmark = True # set False for reproducible results
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)]
print("Using CUDA device0 _CudaDeviceProperties(name='%s', total_memory=%dMB)" % cuda_str = 'Using CUDA' + apex_str
(x[0].name, x[0].total_memory / c)) print("%sdevice0 _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
(cuda_str, x[0].name, x[0].total_memory / c))
if ng > 0: if ng > 0:
# torch.cuda.set_device(0) # OPTIONAL: Set GPU ID # torch.cuda.set_device(0) # OPTIONAL: Set GPU ID
for i in range(1, ng): for i in range(1, ng):
print(" device%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
(i, x[i].name, x[i].total_memory / c)) (' ' * len(cuda_str), i, x[i].name, x[i].total_memory / c))
print('') # skip a line print('') # skip a line
return device return device