updates

train.py

@@ -79,7 +79,7 @@ def train(cfg,
     weights = 'weights' + os.sep
     last = weights + 'last.pt'
     best = weights + 'best.pt'
-    device = torch_utils.select_device()
+    device = torch_utils.select_device(apex=mixed_precision)
     multi_scale = opt.multi_scale
 
     if multi_scale:
@@ -162,7 +162,6 @@ def train(cfg,
     # Mixed precision training https://github.com/NVIDIA/apex
     if mixed_precision:
         model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)
-        print('Using Apex')
 
     # Initialize distributed training
     if torch.cuda.device_count() > 1:

utils/torch_utils.py

@@ -9,7 +9,8 @@ def init_seeds(seed=0):
     # 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()
     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
         ng = torch.cuda.device_count()
         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):
             if i == 1:
                 # torch.cuda.set_device(0)  # OPTIONAL: Set GPU ID