Misc changes for pytorch-translate

fairseq/data/dictionary.py

@@ -106,7 +106,7 @@ class Dictionary(object):
                 multiple of 8, which is important on some hardware (e.g., Nvidia
                 Tensor Cores).
         """
-        if nwords == -1:
+        if nwords <= 0:
             nwords = len(self)
 
         new_indices = dict(zip(self.symbols[:self.nspecial], range(self.nspecial)))
@@ -133,7 +133,7 @@ class Dictionary(object):
                 i += 1
                 threshold_nwords += 1
 
-        assert min(new_count[self.nspecial:]) >= threshold
+        assert len(new_count) == self.nspecial or min(new_count[self.nspecial:]) >= threshold
         assert len(new_symbols) % padding_factor == 0
         assert len(new_symbols) == len(new_indices)
 
@@ -187,12 +187,12 @@ class Dictionary(object):
             d.count.append(count)
         return d
 
-    def save(self, f, threshold=3, nwords=-1):
+    def save(self, f):
         """Stores dictionary into a text file"""
         if isinstance(f, str):
             os.makedirs(os.path.dirname(f), exist_ok=True)
             with open(f, 'w', encoding='utf-8') as fd:
-                return self.save(fd, threshold, nwords)
+                return self.save(fd)
         for symbol, count in zip(self.symbols[self.nspecial:], self.count[self.nspecial:]):
             print('{} {}'.format(symbol, count), file=f)
 

fairseq/data/indexed_dataset.py

@@ -52,8 +52,9 @@ def data_file_path(prefix_path):
 class IndexedDataset(torch.utils.data.Dataset):
     """Loader for TorchNet IndexedDataset"""
 
-    def __init__(self, path):
+    def __init__(self, path, fix_lua_indexing=False):
         super().__init__()
+        self.fix_lua_indexing = fix_lua_indexing
         with open(index_file_path(path), 'rb') as f:
             magic = f.read(8)
             assert magic == b'TNTIDX\x00\x00'
@@ -83,7 +84,10 @@ class IndexedDataset(torch.utils.data.Dataset):
         a = np.empty(tensor_size, dtype=self.dtype)
         self.data_file.seek(self.data_offsets[i] * self.element_size)
         self.data_file.readinto(a)
-        return torch.from_numpy(a).long() - 1  # subtract 1 for 0-based indexing
+        item = torch.from_numpy(a).long()
+        if self.fix_lua_indexing:
+            item -= 1  # subtract 1 for 0-based indexing
+        return item
 
     def __len__(self):
         return self.size
@@ -104,6 +108,7 @@ class IndexedInMemoryDataset(IndexedDataset):
         self.buffer = np.empty(self.data_offsets[-1], dtype=self.dtype)
         self.data_file.readinto(self.buffer)
         self.data_file.close()
+        if self.fix_lua_indexing:
             self.buffer -= 1  # subtract 1 for 0-based indexing
 
     def __del__(self):

fairseq/fp16_trainer.py

@@ -73,7 +73,7 @@ class FP16Trainer(Trainer):
         self.fp32_params.grad = self.fp32_params.data.new(total_param_size)
 
         # create optimizer using the copied FP32 params
-        self.optimizer = optim.build_optimizer(self.args, [self.fp32_params])
+        self._optimizer = optim.build_optimizer(self.args, [self.fp32_params])
         self.lr_scheduler = lr_scheduler.build_lr_scheduler(self.args, self.optimizer)
 
     def save_checkpoint(self, filename, extra_state):

fairseq/optim/lr_scheduler/fixed_schedule.py

@@ -15,6 +15,9 @@ class FixedSchedule(FairseqLRScheduler):
     def __init__(self, args, optimizer):
         super().__init__(args, optimizer)
 
+        # set defaults
+        args.warmup_updates = getattr(args, 'warmup_updates', 0)
+
         self.lr = args.lr[0]
         if args.warmup_updates > 0:
             self.warmup_factor = 1. / args.warmup_updates

fairseq/tasks/language_modeling.py

@@ -50,7 +50,7 @@ class LanguageModelingTask(FairseqTask):
             ds = IndexedRawTextDataset(path, self.dictionary)
             tokens = ds.tokens_list
         elif not self.args.raw_text and IndexedInMemoryDataset.exists(path):
-            ds = IndexedInMemoryDataset(path)
+            ds = IndexedInMemoryDataset(path, fix_lua_indexing=True)
             tokens = ds.buffer
         else:
             raise FileNotFoundError('Dataset not found: {} ({})'.format(split, self.args.data))

fairseq/tasks/translation.py

@@ -89,7 +89,7 @@ class TranslationTask(FairseqTask):
             if self.args.raw_text:
                 return IndexedRawTextDataset(path, dictionary)
             elif IndexedInMemoryDataset.exists(path):
-                return IndexedInMemoryDataset(path)
+                return IndexedInMemoryDataset(path, fix_lua_indexing=True)
             return None
 
         src_dataset = indexed_dataset(prefix + src, self.src_dict)

fairseq/trainer.py

@@ -40,8 +40,6 @@ class Trainer(object):
         self.model = model.cuda()
         self.criterion = criterion.cuda()
 
-        self.optimizer = None
-
         # initialize meters
         self.meters = OrderedDict()
         self.meters['train_loss'] = AverageMeter()
@@ -61,10 +59,17 @@ class Trainer(object):
         self._flat_grads = None
         self._num_updates = 0
         self._optim_history = None
+        self._optimizer = None
+
+    @property
+    def optimizer(self):
+        if self._optimizer is None:
+            self._build_optimizer()
+        return self._optimizer
 
     def _build_optimizer(self):
-        self.optimizer = optim.build_optimizer(self.args, self.model.parameters())
-        self.lr_scheduler = lr_scheduler.build_lr_scheduler(self.args, self.optimizer)
+        self._optimizer = optim.build_optimizer(self.args, self.model.parameters())
+        self.lr_scheduler = lr_scheduler.build_lr_scheduler(self.args, self._optimizer)
 
     def save_checkpoint(self, filename, extra_state):
         """Save all training state in a checkpoint file."""
@@ -93,7 +98,7 @@ class Trainer(object):
 
             self._num_updates = last_optim['num_updates']
 
-        if 'train_meters' in extra_state:
+        if extra_state is not None and 'train_meters' in extra_state:
             self.meters = extra_state['train_meters']
             del extra_state['train_meters']
 
@@ -101,11 +106,6 @@ class Trainer(object):
 
     def train_step(self, sample, update_params=True):
         """Do forward, backward and parameter update."""
-
-        if self.optimizer is None:
-            # initialize optimizer and LR scheduler if hasn't been loaded from the checkpoint
-            self._build_optimizer()
-
         # Set seed based on args.seed and the update number so that we get
         # reproducible results when resuming from checkpoints
         seed = self.args.seed + self.get_num_updates()

fairseq/utils.py

@@ -126,7 +126,7 @@ def _upgrade_state_dict(state):
     if 'train_iterator' not in state['extra_state']:
         state['extra_state']['train_iterator'] = {
             'epoch': state['extra_state']['epoch'],
-            'iterations_in_epoch': 0,
+            'iterations_in_epoch': state['extra_state'].get('batch_offset', 0),
         }
     return state