diff --git a/lr_scheduling.py b/lr_scheduling.py new file mode 100644 index 00000000..6859cb5b --- /dev/null +++ b/lr_scheduling.py @@ -0,0 +1,21 @@ +def poly_lr_scheduler(optimizer, init_lr, iter, lr_decay_iter=1, max_iter=30000, power=0.9,): + """Polynomial decay of learning rate + :param init_lr is base learning rate + :param iter is a current iteration + :param lr_decay_iter how frequently decay occurs, default is 1 + :param max_iter is number of maximum iterations + :param power is a polymomial power + + """ + if iter % lr_decay_iter or iter > max_iter: + return optimizer + + for param_group in optimizer.param_groups: + param_group['lr'] = init_lr*(1 - iter/max_iter)**power + + +def adjust_learning_rate(optimizer, init_lr, epoch): + """Sets the learning rate to the initial LR decayed by 10 every 30 epochs""" + lr = init_lr * (0.1 ** (epoch // 30)) + for param_group in optimizer.param_groups: + param_group['lr'] = lr \ No newline at end of file diff --git a/train.py b/train.py index 03c1b570..29723559 100644 --- a/train.py +++ b/train.py @@ -14,24 +14,7 @@ from ptsemseg.loader import get_loader, get_data_path from ptsemseg.loss import cross_entropy2d from ptsemseg.metrics import scores - - -def poly_lr_scheduler(optimizer, init_lr, iter, lr_decay_iter=1, max_iter=30000, power=0.9,): - """Polynomial decay of learning rate - :param init_lr is base learning rate - :param iter is a current iteration - :param lr_decay_iter how frequently decay occurs, default is 1 - :param max_iter is number of maximum iterations - :param power is a polymomial power - - """ - if iter % lr_decay_iter or iter > max_iter: - return optimizer - - for param_group in optimizer.param_groups: - param_group['lr'] = init_lr*(1 - iter/max_iter)**power - return optimizer - +from lr_scheduling import * def train(args): @@ -74,6 +57,9 @@ def train(args): images = Variable(images) labels = Variable(labels) + iter = len(trainloader)*epoch + i + poly_lr_scheduler(optimizer, args.l_rate, iter) + optimizer.zero_grad() outputs = model(images)