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:28:11 +02:00
parent ce4ea0c332
commit e28437720d
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
train.py
View File

@ -79,7 +79,7 @@ def train(cfg,
weights = 'weights' + os.sep weights = 'weights' + os.sep
last = weights + 'last.pt' last = weights + 'last.pt'
best = weights + 'best.pt' best = weights + 'best.pt'
device = torch_utils.select_device() device = torch_utils.select_device(apex=mixed_precision)
multi_scale = opt.multi_scale multi_scale = opt.multi_scale
if multi_scale: if multi_scale:
@ -162,7 +162,6 @@ 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:

5
utils/torch_utils.py
View File

@ -9,7 +9,8 @@ def init_seeds(seed=0):
# torch.backends.cudnn.deterministic = True # https://pytorch.org/docs/stable/notes/randomness.html # torch.backends.cudnn.deterministic = True # https://pytorch.org/docs/stable/notes/randomness.html
def select_device(force_cpu=False): def select_device(force_cpu=False, apex=False):
# apex if mixed precision training https://github.com/NVIDIA/apex
cuda = False if force_cpu else torch.cuda.is_available() cuda = False if force_cpu else torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu') device = torch.device('cuda:0' if cuda else 'cpu')
@ -20,7 +21,7 @@ def select_device(force_cpu=False):
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 ' cuda_str = 'Using CUDA ' + 'Apex ' if apex else ''
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