Commit 19fafae6 authored by Myle Ott's avatar Myle Ott
Browse files

Allow --lr to specify a fixed learning rate schedule

parent a233fceb

fairseq/multiprocessing_trainer.py

+13 −11
Changes for fairseq/multiprocessing_trainer.py: 13 added lines, 11 removed lines.
Original line number Diff line number Diff line
@@ -68,41 +68,43 @@ class MultiprocessingTrainer(MultiprocessingEventLoop): 
        self.model = model.cuda()

        self.criterion = criterion.cuda()



        # initialize optimizer

        # initialize optimizer and LR scheduler

        self.args.lr = list(map(float, self.args.lr.split(',')))

        self.optimizer = self._build_optimizer()

        self.loss = None



        # initialize LR scheduler

        self.lr_scheduler = self._build_lr_scheduler()



        self.loss = None

        self._max_bsz_seen = 0



    def _build_optimizer(self):

        if self.args.optimizer == 'adagrad':

            return torch.optim.Adagrad(self.model.parameters(), lr=self.args.lr,

            return torch.optim.Adagrad(self.model.parameters(), lr=self.args.lr[0],

                                       weight_decay=self.args.weight_decay)

        elif self.args.optimizer == 'adam':

            return torch.optim.Adam(self.model.parameters(), lr=self.args.lr,

            return torch.optim.Adam(self.model.parameters(), lr=self.args.lr[0],

                                    betas=eval(self.args.adam_betas),

                                    weight_decay=self.args.weight_decay)

        elif self.args.optimizer == 'nag':

            return NAG(self.model.parameters(), lr=self.args.lr,

            return NAG(self.model.parameters(), lr=self.args.lr[0],

                       momentum=self.args.momentum,

                       weight_decay=self.args.weight_decay)

        elif self.args.optimizer == 'sgd':

            return torch.optim.SGD(self.model.parameters(), lr=self.args.lr,

            return torch.optim.SGD(self.model.parameters(), lr=self.args.lr[0],

                                   momentum=self.args.momentum,

                                   weight_decay=self.args.weight_decay)

        else:

            raise ValueError('Unknown optimizer: {}'.format(self.args.optimizer))



    def _build_lr_scheduler(self):

        if self.args.force_anneal > 0:

        if len(self.args.lr) > 1 or self.args.force_anneal > 0:

            lrs = self.args.lr

            def anneal(e):

                if e < self.args.force_anneal:

                    return 1

                    # use fixed LR schedule

                    next_lr = lrs[min(e, len(lrs) - 1)]

                else:

                    return self.args.lrshrink ** (e + 1 - self.args.force_anneal)

                    next_lr = lrs[-1] * self.args.lrshrink ** (e + 1 - self.args.force_anneal)

                return next_lr / lrs[0]  # correct for scaling from LambdaLR

            lr_scheduler = LambdaLR(self.optimizer, anneal)

            lr_scheduler.best = None

        else:

fairseq/options.py

+2 −2
Changes for fairseq/options.py: 2 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -49,8 +49,8 @@ def add_optimization_args(parser): 
    group.add_argument('--optimizer', default='nag', metavar='OPT',

                       choices=MultiprocessingTrainer.OPTIMIZERS,

                       help='optimizer ({})'.format(', '.join(MultiprocessingTrainer.OPTIMIZERS)))

    group.add_argument('--lr', '--learning-rate', default=0.25, type=float, metavar='LR',

                       help='initial learning rate')

    group.add_argument('--lr', '--learning-rate', default='0.25', metavar='LR1,LR2,...,LRn',

                       help='learning rate for the first n epochs with all epochs >n using LRn')

    group.add_argument('--min-lr', metavar='LR', default=1e-5, type=float,

                       help='minimum learning rate')

    group.add_argument('--force-anneal', '--fa', default=0, type=int, metavar='N',