allowing sharded dataset (#696)

Summary:
Co-authored-by: myleott <myleott@fb.com>

Changing `data` to be `str` with colon separated list for loading sharded datasets. This change is useful for loading large datasets that cannot fit into, memory. The large dataset can be sharded and then each shard is loaded in one epoch in roudrobin manner.

For example, if there are `5` shards of data and `10` epochs then the shards will be iterated upon `[0, 1, 2, 3, 4, 0, 1, 2, 3, 4]`.

myleott We need to look into `translation.py` as it currently already expects a list and then concats the datasets.
Pull Request resolved: #696

Differential Revision: D15214049

fbshipit-source-id: 03e43a7b69c7aefada2ca668abf1eac1969fe013

Author: Naman Goyal

fairseq/data/iterators.py

@@ -79,11 +79,12 @@ class EpochBatchIterator(object):
         num_workers (int, optional): how many subprocesses to use for data
             loading. 0 means the data will be loaded in the main process
             (default: 0).
+        epoch (int, optional): The epoch to start the iterator from.
     """
 
     def __init__(
         self, dataset, collate_fn, batch_sampler, seed=1, num_shards=1, shard_id=0,
-        num_workers=0,
+        num_workers=0, epoch=0,
     ):
         assert isinstance(dataset, torch.utils.data.Dataset)
         self.dataset = dataset
@@ -94,7 +95,7 @@ def __init__(
         self.shard_id = shard_id
         self.num_workers = num_workers
 
-        self.epoch = 0
+        self.epoch = epoch
         self._cur_epoch_itr = None
         self._next_epoch_itr = None
         self._supports_prefetch = getattr(dataset, 'supports_prefetch', False)

fairseq/tasks/cross_lingual_lm.py

@@ -42,7 +42,8 @@ class CrossLingualLMTask(FairseqTask):
     @staticmethod
     def add_args(parser):
         """Add task-specific arguments to the parser."""
-        parser.add_argument('data', help='path to data directory')
+        parser.add_argument('data', help='colon separated path to data directories list, \
+                            will be iterated upon during epochs in round-robin manner')
         parser.add_argument('--tokens-per-sample', default=512, type=int,
                             help='max number of total tokens over all segments'
                                  ' per sample')
@@ -106,12 +107,16 @@ def setup_task(cls, args, **kwargs):
 
         return cls(args, dictionary)
 
-    def _load_single_lang_dataset(self, split):
+    def _load_single_lang_dataset(self, split, epoch):
         loaded_datasets = []
 
+        paths = self.args.data.split(':')
+        assert len(paths) > 0
+        data_path = paths[epoch % len(paths)]
+
         for k in itertools.count():
             split_k = split + (str(k) if k > 0 else '')
-            path = os.path.join(self.args.data, split_k)
+            path = os.path.join(data_path, split_k)
 
             if self.args.raw_text and IndexedRawTextDataset.exists(path):
                 ds = IndexedRawTextDataset(path, self.dictionary)
@@ -124,7 +129,7 @@ def _load_single_lang_dataset(self, split):
                 if k > 0:
                     break
                 else:
-                    raise FileNotFoundError('Dataset not found: {} ({})'.format(split, self.args.data))
+                    raise FileNotFoundError('Dataset not found: {} ({})'.format(split, data_path))
 
             # Since we append each block with the classification_token,
             # we need to effectively create blocks of length
@@ -136,7 +141,7 @@ def _load_single_lang_dataset(self, split):
                 )
             )
 
-            print('| {} {} {} examples'.format(self.args.data, split_k, len(loaded_datasets[-1])))
+            print('| {} {} {} examples'.format(data_path, split_k, len(loaded_datasets[-1])))
 
         if len(loaded_datasets) == 1:
             dataset = loaded_datasets[0]
@@ -147,7 +152,7 @@ def _load_single_lang_dataset(self, split):
 
         return dataset, sizes
 
-    def load_dataset(self, split, combine=False, **kwargs):
+    def load_dataset(self, split, epoch=0, combine=False, **kwargs):
         """Load a given dataset split.
         Args:
             split (str): name of the split (e.g., train, valid, test)
@@ -162,7 +167,7 @@ def load_dataset(self, split, combine=False, **kwargs):
             # Datasets are expected to be in "split.lang" format (Eg: train.en)
             language_split = '{}.{}'.format(split, lang)
 
-            block_dataset, sizes = self._load_single_lang_dataset(split=language_split)
+            block_dataset, sizes = self._load_single_lang_dataset(split=language_split, epoch=epoch)
 
             dataset_map[lang] = MaskedLMDataset(
                 dataset=block_dataset,
@@ -182,6 +187,6 @@ def load_dataset(self, split, combine=False, **kwargs):
             dataset_map, default_key=self.default_key
         )
         print('| {} {} {} examples'.format(
-            self.args.data, split, len(self.datasets[split])
+            self.args.data.split(':')[epoch], split, len(self.datasets[split])
             )
-        )
+        )

fairseq/tasks/fairseq_task.py

@@ -92,7 +92,7 @@ def dataset(self, split):
     def get_batch_iterator(
         self, dataset, max_tokens=None, max_sentences=None, max_positions=None,
         ignore_invalid_inputs=False, required_batch_size_multiple=1,
-        seed=1, num_shards=1, shard_id=0, num_workers=0,
+        seed=1, num_shards=1, shard_id=0, num_workers=0, epoch=0,
     ):
         """
         Get an iterator that yields batches of data from the given dataset.
@@ -118,6 +118,7 @@ def get_batch_iterator(
             num_workers (int, optional): how many subprocesses to use for data
                 loading. 0 means the data will be loaded in the main process
                 (default: 0).
+            epoch (int, optional): The epoch to start the iterator from.
 
         Returns:
             ~fairseq.iterators.EpochBatchIterator: a batched iterator over the
@@ -149,6 +150,7 @@ def get_batch_iterator(
             num_shards=num_shards,
             shard_id=shard_id,
             num_workers=num_workers,
+            epoch=epoch,
         )
 
     def build_model(self, args):

fairseq/tasks/language_modeling.py

@@ -104,7 +104,9 @@ def setup_task(cls, args, **kwargs):
         dictionary = None
         output_dictionary = None
         if args.data:
-            dictionary = Dictionary.load(os.path.join(args.data, 'dict.txt'))
+            paths = args.data.split(':')
+            assert len(paths) > 0
+            dictionary = Dictionary.load(os.path.join(paths[0], 'dict.txt'))
             print('| dictionary: {} types'.format(len(dictionary)))
             output_dictionary = dictionary
             if args.output_dictionary_size >= 0:
@@ -136,7 +138,7 @@ def build_model(self, args):
 
         return model
 
-    def load_dataset(self, split, combine=False, **kwargs):
+    def load_dataset(self, split, epoch=0, combine=False, **kwargs):
         """Load a given dataset split.
 
         Args:
@@ -145,9 +147,13 @@ def load_dataset(self, split, combine=False, **kwargs):
 
         loaded_datasets = []
 
+        paths = self.args.data.split(':')
+        assert len(paths) > 0
+        data_path = paths[epoch % len(paths)]
+
         for k in itertools.count():
             split_k = split + (str(k) if k > 0 else '')
-            path = os.path.join(self.args.data, split_k)
+            path = os.path.join(data_path, split_k)
 
             if self.args.raw_text and IndexedRawTextDataset.exists(path):
                 ds = IndexedRawTextDataset(path, self.dictionary)
@@ -160,7 +166,7 @@ def load_dataset(self, split, combine=False, **kwargs):
                 if k > 0:
                     break
                 else:
-                    raise FileNotFoundError('Dataset not found: {} ({})'.format(split, self.args.data))
+                    raise FileNotFoundError('Dataset not found: {} ({})'.format(split, data_path))
 
             loaded_datasets.append(
                 TokenBlockDataset(
@@ -170,7 +176,7 @@ def load_dataset(self, split, combine=False, **kwargs):
                 )
             )
 
-            print('| {} {} {} examples'.format(self.args.data, split_k, len(loaded_datasets[-1])))
+            print('| {} {} {} examples'.format(data_path, split_k, len(loaded_datasets[-1])))
 
             if not combine:
                 break

fairseq/tasks/multilingual_translation.py

@@ -135,7 +135,9 @@ def prepare(cls, args, **kargs):
         # load dictionaries
         dicts = OrderedDict()
         for lang in sorted_langs:
-            dicts[lang] = Dictionary.load(os.path.join(args.data, 'dict.{}.txt'.format(lang)))
+            paths = args.data.split(':')
+            assert len(paths) > 0
+            dicts[lang] = Dictionary.load(os.path.join(paths[0], 'dict.{}.txt'.format(lang)))
             if len(dicts) > 0:
                 assert dicts[lang].pad() == dicts[sorted_langs[0]].pad()
                 assert dicts[lang].eos() == dicts[sorted_langs[0]].eos()
@@ -185,11 +187,15 @@ def alter_dataset_langtok(self, lang_pair_dataset,
             new_tgt_bos=new_tgt_bos,
         )
 
-    def load_dataset(self, split, **kwargs):
+    def load_dataset(self, split, epoch=0, **kwargs):
         """Load a dataset split."""
 
+        paths = self.args.data.split(':')
+        assert len(paths) > 0
+        data_path = paths[epoch % len(paths)]
+
         def split_exists(split, src, tgt, lang):
-            filename = os.path.join(self.args.data, '{}.{}-{}.{}'.format(split, src, tgt, lang))
+            filename = os.path.join(data_path, '{}.{}-{}.{}'.format(split, src, tgt, lang))
             if self.args.raw_text and IndexedRawTextDataset.exists(filename):
                 return True
             elif not self.args.raw_text and IndexedDataset.exists(filename):
@@ -210,17 +216,17 @@ def indexed_dataset(path, dictionary):
         for lang_pair in self.args.lang_pairs:
             src, tgt = lang_pair.split('-')
             if split_exists(split, src, tgt, src):
-                prefix = os.path.join(self.args.data, '{}.{}-{}.'.format(split, src, tgt))
+                prefix = os.path.join(data_path, '{}.{}-{}.'.format(split, src, tgt))
             elif split_exists(split, tgt, src, src):
-                prefix = os.path.join(self.args.data, '{}.{}-{}.'.format(split, tgt, src))
+                prefix = os.path.join(data_path, '{}.{}-{}.'.format(split, tgt, src))
             else:
                 continue
             src_datasets[lang_pair] = indexed_dataset(prefix + src, self.dicts[src])
             tgt_datasets[lang_pair] = indexed_dataset(prefix + tgt, self.dicts[tgt])
-            print('| {} {} {} examples'.format(self.args.data, split, len(src_datasets[lang_pair])))
+            print('| {} {} {} examples'.format(data_path, split, len(src_datasets[lang_pair])))
 
         if len(src_datasets) == 0:
-            raise FileNotFoundError('Dataset not found: {} ({})'.format(split, self.args.data))
+            raise FileNotFoundError('Dataset not found: {} ({})'.format(split, data_path))
 
         def language_pair_dataset(lang_pair):
             src, tgt = lang_pair.split('-')

fairseq/tasks/semisupervised_translation.py

@@ -132,14 +132,18 @@ def setup_task(cls, args, **kwargs):
         dicts, training = MultilingualTranslationTask.prepare(args, **kwargs)
         return cls(args, dicts, training)
 
-    def load_dataset(self, split, **kwargs):
+    def load_dataset(self, split, epoch=0, **kwargs):
         """Load a dataset split."""
 
+        paths = self.args.data.split(':')
+        assert len(paths) > 0
+        data_path = paths[epoch % len(paths)]
+
         def split_exists(split, src, tgt, lang):
             if src is not None:
-                filename = os.path.join(self.args.data, '{}.{}-{}.{}'.format(split, src, tgt, lang))
+                filename = os.path.join(data_path, '{}.{}-{}.{}'.format(split, src, tgt, lang))
             else:
-                filename = os.path.join(self.args.data, '{}.{}-None.{}'.format(split, src, tgt))
+                filename = os.path.join(data_path, '{}.{}-None.{}'.format(split, src, tgt))
             if self.args.raw_text and IndexedRawTextDataset.exists(filename):
                 return True
             elif not self.args.raw_text and IndexedDataset.exists(filename):
@@ -162,25 +166,25 @@ def indexed_dataset(path, dictionary):
             for lang_pair in self.args.lang_pairs:
                 src, tgt = lang_pair.split('-')
                 if split_exists(split, src, tgt, src):
-                    prefix = os.path.join(self.args.data, '{}.{}-{}.'.format(split, src, tgt))
+                    prefix = os.path.join(data_path, '{}.{}-{}.'.format(split, src, tgt))
                 elif split_exists(split, tgt, src, src):
-                    prefix = os.path.join(self.args.data, '{}.{}-{}.'.format(split, tgt, src))
+                    prefix = os.path.join(data_path, '{}.{}-{}.'.format(split, tgt, src))
                 else:
                     continue
                 src_datasets[lang_pair] = indexed_dataset(prefix + src, self.dicts[src])
                 tgt_datasets[lang_pair] = indexed_dataset(prefix + tgt, self.dicts[tgt])
-                print('| parallel-{} {} {} examples'.format(self.args.data, split, len(src_datasets[lang_pair])))
+                print('| parallel-{} {} {} examples'.format(data_path, split, len(src_datasets[lang_pair])))
             if len(src_datasets) == 0:
-                raise FileNotFoundError('Dataset not found: {} ({})'.format(split, self.args.data))
+                raise FileNotFoundError('Dataset not found: {} ({})'.format(split, data_path))
 
         # back translation datasets
         backtranslate_datasets = {}
         if (self.lambda_otf_bt > 0.0 or self.lambda_otf_bt_steps is not None) and split.startswith("train"):
             for lang_pair in self.args.lang_pairs:
                 src, tgt = lang_pair.split('-')
                 if not split_exists(split, tgt, None, tgt):
-                    raise FileNotFoundError('Dataset not found: backtranslation {} ({})'.format(split, self.args.data))
-                filename = os.path.join(self.args.data, '{}.{}-None.{}'.format(split, tgt, tgt))
+                    raise FileNotFoundError('Dataset not found: backtranslation {} ({})'.format(split, data_path))
+                filename = os.path.join(data_path, '{}.{}-None.{}'.format(split, tgt, tgt))
                 dataset = indexed_dataset(filename, self.dicts[tgt])
                 lang_pair_dataset_tgt = LanguagePairDataset(
                     dataset,
@@ -216,7 +220,7 @@ def indexed_dataset(path, dictionary):
                     ).collater,
                 )
                 print('| backtranslate-{}: {} {} {} examples'.format(
-                    tgt, self.args.data, split, len(backtranslate_datasets[lang_pair]),
+                    tgt, data_path, split, len(backtranslate_datasets[lang_pair]),
                 ))
                 self.backtranslate_datasets[lang_pair] = backtranslate_datasets[lang_pair]
 
@@ -227,7 +231,7 @@ def indexed_dataset(path, dictionary):
                 _, tgt = lang_pair.split('-')
                 if not split_exists(split, tgt, None, tgt):
                     continue
-                filename = os.path.join(self.args.data, '{}.{}-None.{}'.format(split, tgt, tgt))
+                filename = os.path.join(data_path, '{}.{}-None.{}'.format(split, tgt, tgt))
                 tgt_dataset1 = indexed_dataset(filename, self.dicts[tgt])
                 tgt_dataset2 = indexed_dataset(filename, self.dicts[tgt])
                 noising_dataset = NoisingDataset(
@@ -255,7 +259,7 @@ def indexed_dataset(path, dictionary):
                     tgt_lang=tgt,
                 )
                 print('| denoising-{}: {} {} {} examples'.format(
-                    tgt, self.args.data, split, len(noising_datasets[lang_pair]),
+                    tgt, data_path, split, len(noising_datasets[lang_pair]),
                 ))
 
         def language_pair_dataset(lang_pair):