Further generalize EpochBatchIterator and move iterators into new file (0a7f9e64) · Commits · Simon Will / fairseq

fairseq/data/init.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -12,4 +12,4 @@ from .language_pair_dataset import LanguagePairDataset
		from .monolingual_dataset import MonolingualDataset
		from .token_block_dataset import TokenBlockDataset

		from .data_utils import EpochBatchIterator
		from .iterators import EpochBatchIterator

fairseq/data/data_utils.py

+0 −153

Original line number	Diff line number	Diff line
		@@ -6,11 +6,9 @@
		# can be found in the PATENTS file in the same directory.

		import contextlib
		import itertools
		import os

		import numpy as np
		import torch


		def infer_language_pair(path):
		@@ -23,60 +21,6 @@ def infer_language_pair(path):
		return src, dst


		class ShardedIterator(object):
		"""A sharded wrapper around an iterable (padded to length)."""

		def __init__(self, iterable, num_shards, shard_id, fill_value=None):
		if shard_id < 0 or shard_id >= num_shards:
		raise ValueError('shard_id must be between 0 and num_shards')

		self._sharded_len = len(iterable) // num_shards
		if len(iterable) % num_shards > 0:
		self._sharded_len += 1

		self.itr = itertools.zip_longest(
		range(self._sharded_len),
		itertools.islice(iterable, shard_id, len(iterable), num_shards),
		fillvalue=fill_value,
		)

		def __len__(self):
		return self._sharded_len

		def __iter__(self):
		return self

		def __next__(self):
		return next(self.itr)[1]


		class CountingIterator(object):
		"""Wrapper around an iterable that maintains the iteration count."""

		def __init__(self, iterable):
		self.iterable = iterable
		self.count = 0
		self.itr = iter(self)

		def __len__(self):
		return len(self.iterable)

		def __iter__(self):
		for x in self.iterable:
		self.count += 1
		yield x

		def __next__(self):
		return next(self.itr)

		def has_next(self):
		return self.count < len(self)

		def skip(self, num_to_skip):
		next(itertools.islice(self.itr, num_to_skip, num_to_skip), None)
		return self


		def collate_tokens(values, pad_idx, eos_idx, left_pad, move_eos_to_beginning=False):
		"""Convert a list of 1d tensors into a padded 2d tensor."""
		size = max(v.size(0) for v in values)
		@@ -96,103 +40,6 @@ def collate_tokens(values, pad_idx, eos_idx, left_pad, move_eos_to_beginning=Fal
		return res


		class EpochBatchIterator(object):
		"""A multi-epoch iterator over a :class:`~torch.utils.data.Dataset`.

		Compared to :class:`~torch.utils.data.DataLoader`, this iterator:

		- can be reused across multiple epochs with the :func:`next_epoch_itr`
		method (optionally shuffled between epochs)
		- can be serialized/deserialized with the :func:`state_dict` and
		:func:`load_state_dict` methods
		- supports sharding with the ``num_shards`` and ``shard_id`` arguments

		Args:
		dataset (Dataset): dataset from which to load the data
		batch_sampler (Sampler): an iterator over batches of indices
		seed (int, optional): seed for random number generator for
		reproducibility. Default: ``1``
		num_shards (int, optional): shard the data iterator into N
		shards. Default: ``1``
		shard_id (int, optional): which shard of the data iterator to
		return. Default: ``0``
		"""

		def __init__(self, dataset, batch_sampler, seed=1, num_shards=1, shard_id=0):
		assert isinstance(dataset, torch.utils.data.Dataset)
		self.dataset = dataset
		self.frozen_batches = tuple(batch_sampler)
		self.seed = seed
		self.num_shards = num_shards
		self.shard_id = shard_id

		self.epoch = 0
		self._cur_epoch_itr = None
		self._next_epoch_itr = None

		def __len__(self):
		return len(self.frozen_batches)

		def next_epoch_itr(self, shuffle=True):
		"""
		Return a new iterator over the dataset.

		Args:
		shuffle (bool, optional): shuffle batches before returning the
		iterator. Default: ``True``
		"""
		if self._next_epoch_itr is not None:
		self._cur_epoch_itr = self._next_epoch_itr
		self._next_epoch_itr = None
		else:
		self.epoch += 1
		self._cur_epoch_itr = self._get_iterator_for_epoch(self.epoch, shuffle)
		return self._cur_epoch_itr

		def end_of_epoch(self):
		"""Returns whether the most recent epoch iterator has been exhausted"""
		return not self._cur_epoch_itr.has_next()

		@property
		def iterations_in_epoch(self):
		"""The number of consumed batches in the current epoch."""
		if self._cur_epoch_itr is not None:
		return self._cur_epoch_itr.count
		elif self._next_epoch_itr is not None:
		return self._next_epoch_itr.count
		return 0

		def state_dict(self):
		return {
		'epoch': self.epoch,
		'iterations_in_epoch': self.iterations_in_epoch,
		}

		def load_state_dict(self, state_dict):
		self.epoch = state_dict['epoch']
		itr_pos = state_dict.get('iterations_in_epoch', 0)
		if itr_pos > 0:
		# fast-forward epoch iterator
		itr = self._get_iterator_for_epoch(self.epoch, state_dict.get('shuffle', True))
		if itr_pos < len(itr):
		self._next_epoch_itr = itr.skip(itr_pos)

		def _get_iterator_for_epoch(self, epoch, shuffle):
		if shuffle:
		# set seed based on the seed and epoch number so that we get
		# reproducible results when resuming from checkpoints
		with numpy_seed(self.seed + epoch):
		batches = list(self.frozen_batches) # copy
		np.random.shuffle(batches)
		else:
		batches = self.frozen_batches
		return CountingIterator(torch.utils.data.DataLoader(
		self.dataset,
		collate_fn=self.dataset.collater,
		batch_sampler=ShardedIterator(batches, self.num_shards, self.shard_id, fill_value=[]),
		))


		@contextlib.contextmanager
		def numpy_seed(seed):
		"""Context manager which seeds the NumPy PRNG with the specified seed and

fairseq/data/iterators.py

0 → 100644

+185 −0

Original line number	Diff line number	Diff line
		# Copyright (c) 2017-present, Facebook, Inc.
		# All rights reserved.
		#
		# This source code is licensed under the license found in the LICENSE file in
		# the root directory of this source tree. An additional grant of patent rights
		# can be found in the PATENTS file in the same directory.

		import itertools

		import numpy as np
		import torch

		from . import data_utils


		class CountingIterator(object):
		"""Wrapper around an iterable that maintains the iteration count.

		Args:
		iterable (iterable): iterable to wrap

		Attributes:
		count (int): number of elements consumed from this iterator
		"""

		def __init__(self, iterable):
		self.iterable = iterable
		self.count = 0
		self.itr = iter(self)

		def __len__(self):
		return len(self.iterable)

		def __iter__(self):
		for x in self.iterable:
		self.count += 1
		yield x

		def __next__(self):
		return next(self.itr)

		def has_next(self):
		"""Whether the iterator has been exhausted."""
		return self.count < len(self)

		def skip(self, num_to_skip):
		"""Fast-forward the iterator by skipping num_to_skip elements."""
		next(itertools.islice(self.itr, num_to_skip, num_to_skip), None)
		return self


		class EpochBatchIterator(object):
		"""A multi-epoch iterator over a :class:`torch.utils.data.Dataset`.

		Compared to :class:`torch.utils.data.DataLoader`, this iterator:

		- can be reused across multiple epochs with the :func:`next_epoch_itr`
		method (optionally shuffled between epochs)
		- can be serialized/deserialized with the :func:`state_dict` and
		:func:`load_state_dict` methods
		- supports sharding with the num_shards and shard_id arguments

		Args:
		dataset (~torch.utils.data.Dataset): dataset from which to load the data
		collate_fn (callable): merges a list of samples to form a mini-batch
		batch_sampler (~torch.utils.data.Sampler): an iterator over batches of
		indices
		seed (int, optional): seed for random number generator for
		reproducibility. Default: ``1``
		num_shards (int, optional): shard the data iterator into N
		shards. Default: ``1``
		shard_id (int, optional): which shard of the data iterator to
		return. Default: ``0``
		"""

		def __init__(self, dataset, collate_fn, batch_sampler, seed=1, num_shards=1, shard_id=0):
		assert isinstance(dataset, torch.utils.data.Dataset)
		self.dataset = dataset
		self.collate_fn = collate_fn
		self.frozen_batches = tuple(batch_sampler)
		self.seed = seed
		self.num_shards = num_shards
		self.shard_id = shard_id

		self.epoch = 0
		self._cur_epoch_itr = None
		self._next_epoch_itr = None

		def __len__(self):
		return len(self.frozen_batches)

		def next_epoch_itr(self, shuffle=True):
		"""Return a new iterator over the dataset.

		Args:
		shuffle (bool, optional): shuffle batches before returning the
		iterator. Default: ``True``
		"""
		if self._next_epoch_itr is not None:
		self._cur_epoch_itr = self._next_epoch_itr
		self._next_epoch_itr = None
		else:
		self.epoch += 1
		self._cur_epoch_itr = self._get_iterator_for_epoch(self.epoch, shuffle)
		return self._cur_epoch_itr

		def end_of_epoch(self):
		"""Returns whether the most recent epoch iterator has been exhausted"""
		return not self._cur_epoch_itr.has_next()

		@property
		def iterations_in_epoch(self):
		"""The number of consumed batches in the current epoch."""
		if self._cur_epoch_itr is not None:
		return self._cur_epoch_itr.count
		elif self._next_epoch_itr is not None:
		return self._next_epoch_itr.count
		return 0

		def state_dict(self):
		"""Returns a dictionary containing a whole state of the iterator."""
		return {
		'epoch': self.epoch,
		'iterations_in_epoch': self.iterations_in_epoch,
		}

		def load_state_dict(self, state_dict):
		"""Copies the state of the iterator from the given state_dict."""
		self.epoch = state_dict['epoch']
		itr_pos = state_dict.get('iterations_in_epoch', 0)
		if itr_pos > 0:
		# fast-forward epoch iterator
		itr = self._get_iterator_for_epoch(self.epoch, state_dict.get('shuffle', True))
		if itr_pos < len(itr):
		self._next_epoch_itr = itr.skip(itr_pos)

		def _get_iterator_for_epoch(self, epoch, shuffle):
		if shuffle:
		# set seed based on the seed and epoch number so that we get
		# reproducible results when resuming from checkpoints
		with data_utils.numpy_seed(self.seed + epoch):
		batches = list(self.frozen_batches) # copy
		np.random.shuffle(batches)
		else:
		batches = self.frozen_batches
		return CountingIterator(torch.utils.data.DataLoader(
		self.dataset,
		collate_fn=self.collate_fn,
		batch_sampler=ShardedIterator(batches, self.num_shards, self.shard_id, fill_value=[]),
		))


		class ShardedIterator(object):
		"""A sharded wrapper around an iterable, padded to length.

		Args:
		iterable (iterable): iterable to wrap
		num_shards (int): number of shards to split the iterable into
		shard_id (int): which shard to iterator over
		fill_value (Any, optional): padding value when the iterable doesn't
		evenly divide num_shards. Default: ``None``
		"""

		def __init__(self, iterable, num_shards, shard_id, fill_value=None):
		if shard_id < 0 or shard_id >= num_shards:
		raise ValueError('shard_id must be between 0 and num_shards')

		self._sharded_len = len(iterable) // num_shards
		if len(iterable) % num_shards > 0:
		self._sharded_len += 1

		self.itr = itertools.zip_longest(
		range(self._sharded_len),
		itertools.islice(iterable, shard_id, len(iterable), num_shards),
		fillvalue=fill_value,
		)

		def __len__(self):
		return self._sharded_len

		def __iter__(self):
		return self

		def __next__(self):
		return next(self.itr)[1]

fairseq/tasks/fairseq_task.py

+3 −2

Original line number	Diff line number	Diff line
		@@ -5,7 +5,7 @@
		# the root directory of this source tree. An additional grant of patent rights
		# can be found in the PATENTS file in the same directory.

		from fairseq.data import data_utils, FairseqDataset
		from fairseq.data import data_utils, FairseqDataset, iterators


		class FairseqTask(object):
		@@ -87,8 +87,9 @@ class FairseqTask(object):
		)

		# return a reusable, sharded iterator
		return data_utils.EpochBatchIterator(
		return iterators.EpochBatchIterator(
		dataset=dataset,
		collate_fn=dataset.collater,
		batch_sampler=batch_sampler,
		seed=seed,
		num_shards=num_shards,

tests/test_data_utils.py→tests/test_iterators.py

+3 −3

Original line number	Diff line number	Diff line
		@@ -7,14 +7,14 @@

		import unittest

		from fairseq.data import data_utils
		from fairseq.data import iterators


		class TestDataUtils(unittest.TestCase):
		class TestIterators(unittest.TestCase):

		def test_counting_iterator(self):
		x = list(range(10))
		itr = data_utils.CountingIterator(x)
		itr = iterators.CountingIterator(x)
		self.assertTrue(itr.has_next())
		self.assertEqual(next(itr), 0)
		self.assertEqual(next(itr), 1)