diff --git a/train.py b/train.py index b1d204f8..98424ed6 100644 --- a/train.py +++ b/train.py @@ -156,7 +156,6 @@ def train(cfg, try: from apex import amp 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 mixed_precision = False diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 3f428fd6..5d08ba20 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -16,17 +16,24 @@ def select_device(force_cpu=False): if not cuda: print('Using CPU') 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 c = 1024 ** 2 # bytes to MB ng = torch.cuda.device_count() x = [torch.cuda.get_device_properties(i) for i in range(ng)] - print("Using CUDA device0 _CudaDeviceProperties(name='%s', total_memory=%dMB)" % - (x[0].name, x[0].total_memory / c)) + cuda_str = 'Using CUDA' + apex_str + print("%sdevice0 _CudaDeviceProperties(name='%s', total_memory=%dMB)" % + (cuda_str, x[0].name, x[0].total_memory / c)) if ng > 0: # torch.cuda.set_device(0) # OPTIONAL: Set GPU ID for i in range(1, ng): - print(" device%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % - (i, x[i].name, x[i].total_memory / c)) + print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % + (' ' * len(cuda_str), i, x[i].name, x[i].total_memory / c)) print('') # skip a line return device