adding a buffered iterator

Summary:
Torch's DataLoader keeps a buffer of only 2 ready batches only, which cannot be changed. This causes a data loading bottleneck at times where data preparation time fluctuates.

Adding BufferedIterator, which is a generic wrapper for an iterator, implementing a buffer using queue.

Adding FairseqTask support, and in BatchSamplerIterator as default.

Reviewed By: myleott

Differential Revision: D21261026

fbshipit-source-id: 23d4bc6181fe1f9a7ee7ad7d18491594725c0f53

Author: Gil Keren

fairseq/data/iterators.py

@@ -6,12 +6,21 @@
 import itertools
 import math
 import os
-
+import time
 import numpy as np
 import torch
-
+import queue
+import logging
+from threading import Thread
 from . import data_utils
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+# Object used by _background_consumer to signal the source is exhausted
+# to the main thread.
+_sentinel = object()
+
 
 class CountingIterator(object):
     """Wrapper around an iterable that maintains the iteration count.
@@ -178,11 +187,14 @@ class EpochBatchIterator(EpochBatchIterating):
             (default: 0).
         epoch (int, optional): the epoch to start the iterator from
             (default: 1).
+        buffer_size (int, optional): the number of batches to keep ready in the
+            queue. Helps speeding up dataloading. When buffer_size is zero, the
+            default torch.utils.data.DataLoader preloading is used.
     """
 
     def __init__(
         self, dataset, collate_fn, batch_sampler, seed=1, num_shards=1, shard_id=0,
-        num_workers=0, epoch=1,
+        num_workers=0, epoch=1, buffer_size=0
     ):
         assert isinstance(dataset, torch.utils.data.Dataset)
         self.dataset = dataset
@@ -192,6 +204,7 @@ def __init__(
         self.num_shards = num_shards
         self.shard_id = shard_id
         self.num_workers = num_workers
+        self.buffer_size = buffer_size
 
         self.epoch = max(epoch, 1)  # we use 1-based indexing for epochs
         self.shuffle = True
@@ -307,16 +320,22 @@ def shuffle_batches(batches, seed):
         if self.num_workers > 0:
             os.environ['PYTHONWARNINGS'] = 'ignore:semaphore_tracker:UserWarning'
 
-        return CountingIterator(
-            torch.utils.data.DataLoader(
-                self.dataset,
-                collate_fn=self.collate_fn,
-                batch_sampler=batches[offset:],
-                num_workers=self.num_workers,
-            ),
-            start=offset,
+        # Create data loader
+        itr = torch.utils.data.DataLoader(
+            self.dataset,
+            collate_fn=self.collate_fn,
+            batch_sampler=batches[offset:],
+            num_workers=self.num_workers,
         )
 
+        # Wrap with a BufferedIterator if needed
+        if self.buffer_size > 0:
+            itr = BufferedIterator(self.buffer_size, itr)
+
+        # Wrap with CoutingIterator
+        itr = CountingIterator(itr, start=offset)
+        return itr
+
 
 class GroupedIterator(object):
     """Wrapper around an iterable that returns groups (chunks) of items.
@@ -382,3 +401,55 @@ def __iter__(self):
 
     def __next__(self):
         return next(self.itr)[1]
+
+
+class BackgroundConsumer(Thread):
+    def __init__(self, queue, source):
+        Thread.__init__(self)
+
+        self._queue = queue
+        self._source = source
+
+    def run(self):
+        for item in self._source:
+            self._queue.put(item)
+
+        # Signal the consumer we are done.
+        self._queue.put(_sentinel)
+
+
+class BufferedIterator(object):
+    def __init__(self, size, iterable):
+        self._queue = queue.Queue(size)
+        self._iterable = iterable
+
+        self._consumer = BackgroundConsumer(self._queue, iterable)
+        self._consumer.daemon = True
+        self._consumer.start()
+
+        self.start_time = time.time()
+        self.warning_time = None
+
+    def __iter__(self):
+        return self
+
+    def __len__(self):
+        return len(self._iterable)
+
+    def __next__(self):
+        # Notify the user if there is a data loading bottleneck
+        if self._queue.qsize() < 2:
+            if time.time() - self.start_time > 5 * 60:
+                if self.warning_time is None or time.time() - self.warning_time > 15 * 60:
+                    logger.info(
+                        "Data loading buffer is empty or nearly empty. This may "
+                        "indicate a data loading bottleneck, and increasing the "
+                        "number of workers may help."
+                    )
+                    self.warning_time = time.time()
+
+        # Get next example
+        item = self._queue.get(True)
+        if item is _sentinel:
+            raise StopIteration()
+        return item

fairseq/options.py

@@ -317,6 +317,8 @@ def add_dataset_args(parser, train=False, gen=False):
     parser.add_argument('--dataset-impl', metavar='FORMAT',
                         choices=get_available_dataset_impl(),
                         help='output dataset implementation')
+    group.add_argument('--data-buffer-size', default=0, type=int, metavar='N',
+                        help='Number of batches to preload')
     if train:
         group.add_argument('--train-subset', default='train', metavar='SPLIT',
                            help='data subset to use for training (e.g. train, valid, test)')

fairseq/tasks/fairseq_task.py

@@ -117,6 +117,7 @@ def get_batch_iterator(
         shard_id=0,
         num_workers=0,
         epoch=1,
+        buffer_size=0
     ):
         """
         Get an iterator that yields batches of data from the given dataset.
@@ -191,6 +192,7 @@ def get_batch_iterator(
             shard_id=shard_id,
             num_workers=num_workers,
             epoch=epoch,
+            buffer_size=buffer_size,
         )
         self.dataset_to_epoch_iter[dataset] = epoch_iter
         return epoch_iter

fairseq/trainer.py

@@ -294,6 +294,7 @@ def get_train_iterator(
             shard_id=self.data_parallel_rank if shard_batch_itr else 0,
             num_workers=self.args.num_workers,
             epoch=epoch,
+            buffer_size=self.args.data_buffer_size,
         )
 
     def get_valid_iterator(
@@ -315,6 +316,7 @@ def get_valid_iterator(
             num_shards=self.data_parallel_world_size,
             shard_id=self.data_parallel_rank,
             num_workers=self.args.num_workers,
+            buffer_size=self.args.data_buffer_size,
         )
 
     def begin_epoch(self, epoch):