diff --git a/train.py b/train.py index be7ffb3d..663f399d 100644 --- a/train.py +++ b/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: diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 6dabe6ae..8219ccc9 100644 --- a/utils/torch_utils.py +++ b/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