From 806d7b92d88364f2873563411078019e9b3b05c8 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 11 Sep 2019 14:25:48 +0200 Subject: [PATCH] updates --- train.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/train.py b/train.py index c2bb0c33..37e4ab13 100644 --- a/train.py +++ b/train.py @@ -1,5 +1,4 @@ import argparse -import time import torch.distributed as dist import torch.optim as optim @@ -88,9 +87,11 @@ def train(): else: pg0 += [v] # parameter group 0 - # optimizer = optim.Adam(pg0, lr=hyp['lr0']) - # optimizer = AdaBound(pg0, lr=hyp['lr0'], final_lr=0.1) - optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) + if opt.adam: + optimizer = optim.Adam(pg0, lr=hyp['lr0']) + # optimizer = AdaBound(pg0, lr=hyp['lr0'], final_lr=0.1) + else: + optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay del pg0, pg1 @@ -388,6 +389,7 @@ if __name__ == '__main__': parser.add_argument('--prebias', action='store_true', help='transfer-learn yolo biases prior to training') parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied') parser.add_argument('--device', default='', help='select device if multi-gpu, i.e. 0 or 0,1') + parser.add_argument('--adam', action='store_true', help='use adam optimizer') parser.add_argument('--var', type=float, help='debug variable') opt = parser.parse_args() opt.weights = 'weights/last.pt' if opt.resume else opt.weights