Allow schedule for update-freq

fairseq/options.py

@@ -58,6 +58,7 @@ def parse_args_and_arch(parser, input_args=None):
 
     # Post-process args.
     args.lr = list(map(float, args.lr.split(',')))
+    args.update_freq = list(map(float, args.update_freq.split(',')))
     if args.max_sentences_valid is None:
         args.max_sentences_valid = args.max_sentences
 
@@ -152,6 +153,8 @@ def add_optimization_args(parser):
     group.add_argument('--sentence-avg', action='store_true',
                        help='normalize gradients by the number of sentences in a batch'
                             ' (default is to normalize by number of tokens)')
+    group.add_argument('--update-freq', default='1', metavar='N',
+                       help='update parameters every N_i batches, when in epoch i')
 
     # Optimizer definitions can be found under fairseq/optim/
     group.add_argument('--optimizer', default='nag', metavar='OPT',
@@ -174,8 +177,6 @@ def add_optimization_args(parser):
     group.add_argument('--min-lr', default=1e-5, type=float, metavar='LR',
                        help='minimum learning rate')
 
-    group.add_argument('--update-freq', default=1, type=int, metavar='N',
-                       help='update parameters every N batches')
     return group
 
 

singleprocess_train.py

@@ -132,12 +132,18 @@ def train(args, trainer, itr, epoch):
         if meter is not None:
             meter.reset()
 
+    # update parameters every N batches
+    if epoch <= len(args.update_freq):
+        update_freq = args.update_freq[epoch - 1]
+    else:
+        update_freq = args.update_freq[-1]
+
     extra_meters = collections.defaultdict(lambda: AverageMeter())
     max_update = args.max_update or math.inf
     num_batches = len(itr)
     progress = progress_bar.build_progress_bar(args, itr, epoch, no_progress_bar='simple')
     for i, sample in enumerate(progress):
-        if i < num_batches - 1 and (i + 1) % args.update_freq > 0:
+        if i < num_batches - 1 and (i + 1) % update_freq > 0:
             # buffer updates according to --update-freq
             trainer.train_step(sample, update_params=False)
             continue