From d5815ebfd28ef44deeae0b61c8dc41e18f4197cf Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 27 Feb 2020 13:40:14 -0800 Subject: [PATCH] updates --- train.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/train.py b/train.py index 3a35d65c..c38d375b 100644 --- a/train.py +++ b/train.py @@ -132,6 +132,10 @@ def train(): # possible weights are '*.weights', 'yolov3-tiny.conv.15', 'darknet53.conv.74' etc. load_darknet_weights(model, weights) + # Mixed precision training https://github.com/NVIDIA/apex + if mixed_precision: + model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0) + # Scheduler https://github.com/ultralytics/yolov3/issues/238 # lf = lambda x: 1 - x / epochs # linear ramp to zero # lf = lambda x: 10 ** (hyp['lrf'] * x / epochs) # exp ramp @@ -153,10 +157,6 @@ def train(): # plt.tight_layout() # plt.savefig('LR.png', dpi=300) - # Mixed precision training https://github.com/NVIDIA/apex - if mixed_precision: - model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0) - # Initialize distributed training if device.type != 'cpu' and torch.cuda.device_count() > 1: dist.init_process_group(backend='nccl', # 'distributed backend'