Add BufferedIterator (#419)

Summary:
This improves performance for datasets that load data lazily. Enabled by default since it shouldn't compromise performance for non-lazy datasets.
Pull Request resolved: https://github.com/pytorch/fairseq/pull/419

Differential Revision: D13546585

Pulled By: myleott

fbshipit-source-id: f6152e2047291b0d68cd7506cd772b0caafe95be

fairseq/data/iterators.py

@@ -7,6 +7,8 @@
 
 import itertools
 import math
+import queue
+import threading
 
 import numpy as np
 import torch
@@ -67,14 +69,18 @@ class EpochBatchIterator(object):
         batch_sampler (~torch.utils.data.Sampler): an iterator over batches of
             indices
         seed (int, optional): seed for random number generator for
-            reproducibility. Default: ``1``
+            reproducibility. Default: 1
         num_shards (int, optional): shard the data iterator into N
-            shards. Default: ``1``
+            shards. Default: 1
         shard_id (int, optional): which shard of the data iterator to
-            return. Default: ``0``
+            return. Default: 0
+        buffer_size (int, optional): number of batches to buffer. Default: 5
     """
 
-    def __init__(self, dataset, collate_fn, batch_sampler, seed=1, num_shards=1, shard_id=0):
+    def __init__(
+        self, dataset, collate_fn, batch_sampler, seed=1, num_shards=1, shard_id=0,
+        buffer_size=5,
+    ):
         assert isinstance(dataset, torch.utils.data.Dataset)
         self.dataset = dataset
         self.collate_fn = collate_fn
@@ -82,6 +88,7 @@ class EpochBatchIterator(object):
         self.seed = seed
         self.num_shards = num_shards
         self.shard_id = shard_id
+        self.buffer_size = buffer_size
 
         self.epoch = 0
         self._cur_epoch_itr = None
@@ -172,13 +179,50 @@ class EpochBatchIterator(object):
                 batches = self.frozen_batches
             batches = ShardedIterator(batches, self.num_shards, self.shard_id, fill_value=[])
 
-        return CountingIterator(torch.utils.data.DataLoader(
+        return CountingIterator(BufferedIterator(
+            torch.utils.data.DataLoader(
                 self.dataset,
                 collate_fn=self.collate_fn,
                 batch_sampler=batches,
+            ),
+            buffer_size=self.buffer_size,
         ))
 
 
+class BufferedIterator(object):
+    """Wrapper around an iterable that prefetches items into a buffer.
+
+    Args:
+        iterable (iterable): iterable to wrap
+        buffer_size (int): number of items to prefetch and buffer
+    """
+
+    def __init__(self, iterable, buffer_size):
+        self.iterable = iterable
+
+        self.q = queue.Queue(maxsize=buffer_size)
+        self.thread = threading.Thread(target=self._load_q, daemon=True)
+        self.thread.start()
+
+    def __len__(self):
+        return len(self.iterable)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        x = self.q.get()
+        if x is None:
+            self.thread.join()
+            raise StopIteration
+        return x[0]
+
+    def _load_q(self):
+        for x in self.iterable:
+            self.q.put([x])  # wrap in list so that it's never None
+        self.q.put(None)
+
+
 class GroupedIterator(object):
     """Wrapper around an iterable that returns groups (chunks) of items.