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-24 18:22:59 +02:00
parent 3cfc84a183
commit ce4ea0c332
2 changed files with 2 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
train.py
View File

@ -162,6 +162,7 @@ def train(cfg,
# Mixed precision training https://github.com/NVIDIA/apex # Mixed precision training https://github.com/NVIDIA/apex
if mixed_precision: if mixed_precision:
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')
# Initialize distributed training # Initialize distributed training
if torch.cuda.device_count() > 1: if torch.cuda.device_count() > 1:

8
utils/torch_utils.py
View File

@ -16,17 +16,11 @@ 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 = '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)]
cuda_str = 'Using CUDA ' + apex_str cuda_str = 'Using CUDA '
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 # torch.cuda.set_device(0) # OPTIONAL: Set GPU ID