duplicates

Author: Elad Hoffer

main.py

@@ -101,6 +101,8 @@
                     help='fixed sequence length')
 parser.add_argument('--chunk-batch', default=1, type=int,
                     help='chunk batch size for multiple passes (training) -- used to fit large batches in memory')
+parser.add_argument('--duplicates', default=1, type=int,
+                    help='number of duplicates over singel example')                    
 parser.add_argument('--seed', default=123, type=int,
                     help='random seed (default: 123)')
 
@@ -205,6 +207,7 @@ def main(args):
         keep_checkpoints=args.keep_checkpoints,
         max_tokens=args.max_tokens,
         chunk_batch=args.chunk_batch,
+        duplicates=args.duplicates,
         distributed=args.distributed,
         local_rank=args.local_rank,
         device_ids=args.device_ids,

seq2seq/tools/trainer.py

@@ -51,14 +51,14 @@ def forward(self, module_inputs, target):
             return loss, nll
 
 
-def _chunk_tuple(seq_tuple, num_chunks, batch_first=True):
+def _chunk_tuple(seq_tuple, num_chunks, duplicates=1, batch_first=True):
     if num_chunks == 1:
-        return [seq_tuple]
+        return [seq_tuple] * duplicates
     seq, length = seq_tuple
     batch_dim = 0 if batch_first else 1
     chunked_length = [l.tolist()
                       for l in torch.tensor(length).chunk(num_chunks)]
-    return zip(seq.chunk(num_chunks, dim=batch_dim), chunked_length)
+    return list(zip(seq.chunk(num_chunks, dim=batch_dim), chunked_length)) * duplicates
 
 
 def _batch_max_tokens(src_tuple, target_tuple, max_tokens, batch_first=True, log=True):
@@ -107,6 +107,7 @@ def __init__(self, model, regime=None,
                  embedding_grad_clip=None,
                  max_tokens=None,
                  chunk_batch=1,
+                 duplicates=1,
                  save_info={},
                  save_path='.',
                  checkpoint_filename='checkpoint%s.pth',
@@ -132,6 +133,7 @@ def __init__(self, model, regime=None,
         self.dtype = dtype
         self.max_tokens = max_tokens
         self.chunk_batch = chunk_batch
+        self.duplicates = duplicates
         self.print_freq = print_freq
         self.eval_freq = eval_freq
         self.perplexity = float('inf')
@@ -171,8 +173,8 @@ def iterate(self, src_tuple_batch, target_tuple_batch, training=True, chunk_batc
             self.optimizer.zero_grad()
 
         repacked_inputs = []
-        for src_tuple, target_tuple in zip(_chunk_tuple(src_tuple_batch, chunk_batch, self.batch_first),
-                                           _chunk_tuple(target_tuple_batch, chunk_batch, self.batch_first)):
+        for src_tuple, target_tuple in zip(_chunk_tuple(src_tuple_batch, chunk_batch, self.duplicates, self.batch_first),
+                                           _chunk_tuple(target_tuple_batch, chunk_batch, self.duplicates, self.batch_first)):
             # limit number of tokens to avoid gpu overload
             if training and self.max_tokens is not None:
                 src_tuple, target_tuple = _batch_max_tokens(