wehub-resource-sync
654 lines
22 KiB
Python
654 lines
22 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import itertools
|
|
import logging
|
|
import math
|
|
import operator
|
|
import os
|
|
import queue
|
|
import time
|
|
from threading import Thread
|
|
|
|
import numpy as np
|
|
import torch
|
|
from fairseq.data import data_utils
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 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.
|
|
|
|
Args:
|
|
iterable (iterable): iterable to wrap
|
|
start (int): starting iteration count. Note that this doesn't
|
|
actually advance the iterator.
|
|
total (int): override the iterator length returned by
|
|
``__len__``. This can be used to truncate *iterator*.
|
|
|
|
Attributes:
|
|
n (int): number of elements consumed from this iterator
|
|
"""
|
|
|
|
def __init__(self, iterable, start=None, total=None):
|
|
self.iterable = iterable
|
|
self.itr = iter(self)
|
|
|
|
if start is None:
|
|
self.n = getattr(iterable, "n", 0)
|
|
else:
|
|
self.n = start
|
|
|
|
if total is None:
|
|
self.total = self.n + len(iterable)
|
|
else:
|
|
self.total = total
|
|
|
|
def __len__(self):
|
|
return self.total
|
|
|
|
def __iter__(self):
|
|
for x in self.iterable:
|
|
if self.n >= self.total:
|
|
raise RuntimeError(
|
|
"Mismatch between actual and expected iterable length. "
|
|
"This may be caused by resuming training from a checkpoint using "
|
|
"a different number of GPUs, in which case you can try the "
|
|
"--reset-dataloader option. Alternatively you may have a train or "
|
|
"validation set that is smaller than the number of GPUs. If none "
|
|
"of these apply, please report this to the fairseq developers."
|
|
)
|
|
self.n += 1
|
|
yield x
|
|
|
|
def __next__(self):
|
|
return next(self.itr)
|
|
|
|
def has_next(self):
|
|
"""Whether the iterator has been exhausted."""
|
|
return self.n < 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
|
|
|
|
def take(self, n):
|
|
"""
|
|
Truncates the iterator to n elements at most.
|
|
"""
|
|
self.total = min(self.total, n)
|
|
|
|
# Propagate this change to the underlying iterator
|
|
# Only take after what we have already consumed (i.e. after restarting
|
|
# from checkpoint mid epoch, we have to subtract self.n which is the
|
|
# starting point)
|
|
#
|
|
# This to maintain the invariant self.total = self.n + len(iterable),
|
|
# before calling __next__ or __iter__
|
|
propagated_take = max(n - self.n, 0)
|
|
if hasattr(self.iterable, "take"):
|
|
self.iterable.take(propagated_take)
|
|
else:
|
|
self.iterable = itertools.islice(self.iterable, propagated_take)
|
|
|
|
|
|
class EpochBatchIterating(object):
|
|
def __len__(self) -> int:
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def next_epoch_idx(self):
|
|
raise NotImplementedError
|
|
|
|
def next_epoch_itr(
|
|
self, shuffle=True, fix_batches_to_gpus=False, set_dataset_epoch=True
|
|
):
|
|
"""Return a new iterator over the dataset.
|
|
|
|
Args:
|
|
shuffle (bool, optional): shuffle batches before returning the
|
|
iterator (default: True).
|
|
fix_batches_to_gpus (bool, optional): ensure that batches are always
|
|
allocated to the same shards across epochs. Requires
|
|
that :attr:`dataset` supports prefetching (default: False).
|
|
set_dataset_epoch (bool, optional): update the wrapped Dataset with
|
|
the new epoch number (default: True).
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def end_of_epoch(self) -> bool:
|
|
"""Returns whether the most recent epoch iterator has been exhausted"""
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def iterations_in_epoch(self) -> int:
|
|
"""The number of consumed batches in the current epoch."""
|
|
raise NotImplementedError
|
|
|
|
def state_dict(self):
|
|
"""Returns a dictionary containing a whole state of the iterator."""
|
|
raise NotImplementedError
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Copies the state of the iterator from the given *state_dict*."""
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def first_batch(self):
|
|
return "DUMMY"
|
|
|
|
|
|
class StreamingEpochBatchIterator(EpochBatchIterating):
|
|
"""A steaming-style iterator over a :class:`torch.utils.data.IterableDataset`.
|
|
|
|
Args:
|
|
dataset (~torch.utils.data.Dataset): dataset from which to load the data
|
|
max_sentences: batch size
|
|
collate_fn (callable): merges a list of samples to form a mini-batch
|
|
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
|
|
(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.
|
|
timeout (int, optional): if positive, the timeout value for collecting a batch
|
|
from workers. Should always be non-negative (default: ``0``).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset,
|
|
max_sentences=1,
|
|
collate_fn=None,
|
|
epoch=1,
|
|
num_workers=0,
|
|
buffer_size=0,
|
|
timeout=0,
|
|
):
|
|
assert isinstance(dataset, torch.utils.data.IterableDataset)
|
|
self.dataset = dataset
|
|
self.max_sentences = max_sentences
|
|
self.collate_fn = collate_fn
|
|
self.epoch = max(epoch, 1) # we use 1-based indexing for epochs
|
|
self.num_workers = num_workers
|
|
# This upper limit here is to prevent people from abusing this feature
|
|
# in a shared computing environment.
|
|
self.buffer_size = min(buffer_size, 20)
|
|
self.timeout = timeout
|
|
|
|
self._current_epoch_iterator = None
|
|
|
|
@property
|
|
def next_epoch_idx(self):
|
|
"""Return the epoch index after *next_epoch_itr* is called."""
|
|
if self._current_epoch_iterator is not None and self.end_of_epoch():
|
|
return self.epoch + 1
|
|
else:
|
|
return self.epoch
|
|
|
|
def next_epoch_itr(
|
|
self, shuffle=True, fix_batches_to_gpus=False, set_dataset_epoch=True
|
|
):
|
|
self.epoch = self.next_epoch_idx
|
|
if set_dataset_epoch and hasattr(self.dataset, "set_epoch"):
|
|
self.dataset.set_epoch(self.epoch)
|
|
self._current_epoch_iterator = self._get_iterator_for_epoch(self.epoch, shuffle)
|
|
return self._current_epoch_iterator
|
|
|
|
def end_of_epoch(self) -> bool:
|
|
return not self._current_epoch_iterator.has_next()
|
|
|
|
@property
|
|
def iterations_in_epoch(self) -> int:
|
|
if self._current_epoch_iterator is not None:
|
|
return self._current_epoch_iterator.n
|
|
return 0
|
|
|
|
def state_dict(self):
|
|
return {
|
|
"epoch": self.epoch,
|
|
}
|
|
|
|
def load_state_dict(self, state_dict):
|
|
self.epoch = state_dict["epoch"]
|
|
|
|
def _get_iterator_for_epoch(self, epoch, shuffle, offset=0):
|
|
if self.num_workers > 0:
|
|
os.environ["PYTHONWARNINGS"] = "ignore:semaphore_tracker:UserWarning"
|
|
|
|
# Create data loader
|
|
worker_init_fn = getattr(self.dataset, "worker_init_fn", None)
|
|
itr = torch.utils.data.DataLoader(
|
|
self.dataset,
|
|
batch_size=self.max_sentences,
|
|
collate_fn=self.collate_fn,
|
|
num_workers=self.num_workers,
|
|
timeout=self.timeout,
|
|
worker_init_fn=worker_init_fn,
|
|
)
|
|
|
|
# Wrap with a BufferedIterator if needed
|
|
if self.buffer_size > 0:
|
|
itr = BufferedIterator(self.buffer_size, itr)
|
|
|
|
# Wrap with CountingIterator
|
|
itr = CountingIterator(itr, start=offset)
|
|
|
|
return itr
|
|
|
|
|
|
class EpochBatchIterator(EpochBatchIterating):
|
|
"""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 or a callable): an iterator over batches of
|
|
indices, or a callable to create such an iterator (~torch.utils.data.Sampler).
|
|
A callable batch_sampler will be called for each epoch to enable per epoch dynamic
|
|
batch iterators defined by this callable batch_sampler.
|
|
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).
|
|
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
|
|
(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.
|
|
timeout (int, optional): if positive, the timeout value for collecting a batch
|
|
from workers. Should always be non-negative (default: ``0``).
|
|
disable_shuffling (bool, optional): force disable shuffling
|
|
(default: ``False``).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset,
|
|
collate_fn,
|
|
batch_sampler,
|
|
seed=1,
|
|
num_shards=1,
|
|
shard_id=0,
|
|
num_workers=0,
|
|
epoch=1,
|
|
buffer_size=0,
|
|
timeout=0,
|
|
disable_shuffling=False,
|
|
):
|
|
assert isinstance(dataset, torch.utils.data.Dataset)
|
|
self.dataset = dataset
|
|
self.collate_fn = collate_fn
|
|
self.batch_sampler = batch_sampler
|
|
self._frozen_batches = (
|
|
tuple(batch_sampler) if not callable(batch_sampler) else None
|
|
)
|
|
self.seed = seed
|
|
self.num_shards = num_shards
|
|
self.shard_id = shard_id
|
|
self.num_workers = num_workers
|
|
# This upper limit here is to prevent people from abusing this feature
|
|
# in a shared computing environment.
|
|
self.buffer_size = min(buffer_size, 20)
|
|
self.timeout = timeout
|
|
self.disable_shuffling = disable_shuffling
|
|
|
|
self.epoch = max(epoch, 1) # we use 1-based indexing for epochs
|
|
self.shuffle = not disable_shuffling
|
|
self._cur_epoch_itr = None
|
|
self._next_epoch_itr = None
|
|
self._supports_prefetch = getattr(dataset, "supports_prefetch", False)
|
|
|
|
@property
|
|
def frozen_batches(self):
|
|
if self._frozen_batches is None:
|
|
self._frozen_batches = tuple(self.batch_sampler(self.dataset, self.epoch))
|
|
return self._frozen_batches
|
|
|
|
@property
|
|
def first_batch(self):
|
|
if len(self.frozen_batches) == 0:
|
|
raise Exception(
|
|
"The dataset is empty. This could indicate "
|
|
"that all elements in the dataset have been skipped. "
|
|
"Try increasing the max number of allowed tokens or using "
|
|
"a larger dataset."
|
|
)
|
|
|
|
if getattr(self.dataset, "supports_fetch_outside_dataloader", True):
|
|
return self.collate_fn([self.dataset[i] for i in self.frozen_batches[0]])
|
|
else:
|
|
return "DUMMY"
|
|
|
|
def __len__(self):
|
|
return int(math.ceil(len(self.frozen_batches) / float(self.num_shards)))
|
|
|
|
@property
|
|
def n(self):
|
|
return self.iterations_in_epoch
|
|
|
|
@property
|
|
def next_epoch_idx(self):
|
|
"""Return the epoch index after *next_epoch_itr* is called."""
|
|
if self._next_epoch_itr is not None:
|
|
return self.epoch
|
|
elif self._cur_epoch_itr is not None and self.end_of_epoch():
|
|
return self.epoch + 1
|
|
else:
|
|
return self.epoch
|
|
|
|
def next_epoch_itr(
|
|
self, shuffle=True, fix_batches_to_gpus=False, set_dataset_epoch=True
|
|
):
|
|
"""Return a new iterator over the dataset.
|
|
|
|
Args:
|
|
shuffle (bool, optional): shuffle batches before returning the
|
|
iterator (default: True).
|
|
fix_batches_to_gpus (bool, optional): ensure that batches are always
|
|
allocated to the same shards across epochs. Requires
|
|
that :attr:`dataset` supports prefetching (default: False).
|
|
set_dataset_epoch (bool, optional): update the wrapped Dataset with
|
|
the new epoch number (default: True).
|
|
"""
|
|
if self.disable_shuffling:
|
|
shuffle = False
|
|
self.epoch = self.next_epoch_idx
|
|
if set_dataset_epoch and hasattr(self.dataset, "set_epoch"):
|
|
self.dataset.set_epoch(self.epoch)
|
|
if self._next_epoch_itr is not None:
|
|
self._cur_epoch_itr = self._next_epoch_itr
|
|
self._next_epoch_itr = None
|
|
else:
|
|
if callable(self.batch_sampler):
|
|
# reset _frozen_batches to refresh the next epoch
|
|
self._frozen_batches = None
|
|
self._cur_epoch_itr = self._get_iterator_for_epoch(
|
|
self.epoch,
|
|
shuffle,
|
|
fix_batches_to_gpus=fix_batches_to_gpus,
|
|
)
|
|
self.shuffle = shuffle
|
|
return self._cur_epoch_itr
|
|
|
|
def end_of_epoch(self) -> bool:
|
|
"""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.n
|
|
elif self._next_epoch_itr is not None:
|
|
return self._next_epoch_itr.n
|
|
return 0
|
|
|
|
def state_dict(self):
|
|
"""Returns a dictionary containing a whole state of the iterator."""
|
|
if self.end_of_epoch():
|
|
epoch = self.epoch + 1
|
|
iter_in_epoch = 0
|
|
else:
|
|
epoch = self.epoch
|
|
iter_in_epoch = self.iterations_in_epoch
|
|
return {
|
|
"version": 2,
|
|
"epoch": epoch,
|
|
"iterations_in_epoch": iter_in_epoch,
|
|
"shuffle": self.shuffle,
|
|
}
|
|
|
|
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)
|
|
version = state_dict.get("version", 1)
|
|
if itr_pos > 0:
|
|
# fast-forward epoch iterator
|
|
self._next_epoch_itr = self._get_iterator_for_epoch(
|
|
self.epoch,
|
|
shuffle=state_dict.get("shuffle", True),
|
|
offset=itr_pos,
|
|
)
|
|
if self._next_epoch_itr is None:
|
|
if version == 1:
|
|
# legacy behavior: we finished the epoch, increment epoch counter
|
|
self.epoch += 1
|
|
else:
|
|
raise RuntimeError(
|
|
"Cannot resume training due to dataloader mismatch, please "
|
|
"report this to the fairseq developers. You can relaunch "
|
|
"training with `--reset-dataloader` and it should work."
|
|
)
|
|
else:
|
|
self._next_epoch_itr = None
|
|
|
|
def _get_iterator_for_epoch(
|
|
self, epoch, shuffle, fix_batches_to_gpus=False, offset=0
|
|
):
|
|
def shuffle_batches(batches, seed):
|
|
with data_utils.numpy_seed(seed):
|
|
np.random.shuffle(batches)
|
|
return batches
|
|
|
|
if self._supports_prefetch:
|
|
batches = self.frozen_batches
|
|
|
|
if shuffle and not fix_batches_to_gpus:
|
|
batches = shuffle_batches(list(batches), self.seed + epoch)
|
|
|
|
batches = list(
|
|
ShardedIterator(batches, self.num_shards, self.shard_id, fill_value=[])
|
|
)
|
|
self.dataset.prefetch([i for s in batches for i in s])
|
|
|
|
if shuffle and fix_batches_to_gpus:
|
|
batches = shuffle_batches(batches, self.seed + epoch + self.shard_id)
|
|
else:
|
|
if shuffle:
|
|
batches = shuffle_batches(list(self.frozen_batches), self.seed + epoch)
|
|
else:
|
|
batches = self.frozen_batches
|
|
batches = list(
|
|
ShardedIterator(batches, self.num_shards, self.shard_id, fill_value=[])
|
|
)
|
|
|
|
if offset > 0 and offset >= len(batches):
|
|
return None
|
|
|
|
if self.num_workers > 0:
|
|
os.environ["PYTHONWARNINGS"] = "ignore:semaphore_tracker:UserWarning"
|
|
|
|
# Create data loader
|
|
itr = torch.utils.data.DataLoader(
|
|
self.dataset,
|
|
collate_fn=self.collate_fn,
|
|
batch_sampler=batches[offset:],
|
|
num_workers=self.num_workers,
|
|
timeout=self.timeout,
|
|
)
|
|
|
|
# Wrap with a BufferedIterator if needed
|
|
if self.buffer_size > 0:
|
|
itr = BufferedIterator(self.buffer_size, itr)
|
|
|
|
# Wrap with CountingIterator
|
|
itr = CountingIterator(itr, start=offset)
|
|
return itr
|
|
|
|
|
|
class GroupedIterator(CountingIterator):
|
|
"""Wrapper around an iterable that returns groups (chunks) of items.
|
|
|
|
Args:
|
|
iterable (iterable): iterable to wrap
|
|
chunk_size (int): size of each chunk
|
|
|
|
Attributes:
|
|
n (int): number of elements consumed from this iterator
|
|
"""
|
|
|
|
def __init__(self, iterable, chunk_size):
|
|
itr = _chunk_iterator(iterable, chunk_size)
|
|
super().__init__(
|
|
itr,
|
|
start=int(math.ceil(getattr(iterable, "n", 0) / float(chunk_size))),
|
|
total=int(math.ceil(len(iterable) / float(chunk_size))),
|
|
)
|
|
self.chunk_size = chunk_size
|
|
|
|
|
|
def _chunk_iterator(itr, chunk_size):
|
|
chunk = []
|
|
for x in itr:
|
|
chunk.append(x)
|
|
if len(chunk) == chunk_size:
|
|
yield chunk
|
|
chunk = []
|
|
if len(chunk) > 0:
|
|
yield chunk
|
|
|
|
|
|
class ShardedIterator(CountingIterator):
|
|
"""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).
|
|
|
|
Attributes:
|
|
n (int): number of elements consumed from this iterator
|
|
"""
|
|
|
|
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")
|
|
sharded_len = int(math.ceil(len(iterable) / float(num_shards)))
|
|
itr = map(
|
|
operator.itemgetter(1),
|
|
itertools.zip_longest(
|
|
range(sharded_len),
|
|
itertools.islice(iterable, shard_id, len(iterable), num_shards),
|
|
fillvalue=fill_value,
|
|
),
|
|
)
|
|
super().__init__(
|
|
itr,
|
|
start=int(math.ceil(getattr(iterable, "n", 0) / float(num_shards))),
|
|
total=sharded_len,
|
|
)
|
|
|
|
|
|
class BackgroundConsumer(Thread):
|
|
def __init__(self, queue, source, max_len):
|
|
Thread.__init__(self)
|
|
|
|
self._queue = queue
|
|
self._source = source
|
|
self._max_len = max_len
|
|
self.count = 0
|
|
|
|
def run(self):
|
|
try:
|
|
for item in self._source:
|
|
self._queue.put(item)
|
|
|
|
# Stop if we reached the maximum length
|
|
self.count += 1
|
|
if self._max_len is not None and self.count >= self._max_len:
|
|
break
|
|
|
|
# Signal the consumer we are done.
|
|
self._queue.put(_sentinel)
|
|
except Exception as e:
|
|
self._queue.put(e)
|
|
|
|
|
|
class BufferedIterator(object):
|
|
def __init__(self, size, iterable):
|
|
self._queue = queue.Queue(size)
|
|
self._iterable = iterable
|
|
self._consumer = None
|
|
|
|
self.start_time = time.time()
|
|
self.warning_time = None
|
|
|
|
self.total = len(iterable)
|
|
|
|
def _create_consumer(self):
|
|
self._consumer = BackgroundConsumer(
|
|
self._queue,
|
|
self._iterable,
|
|
self.total,
|
|
)
|
|
self._consumer.daemon = True
|
|
self._consumer.start()
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __len__(self):
|
|
return self.total
|
|
|
|
def take(self, n):
|
|
self.total = min(self.total, n)
|
|
|
|
# Propagate this change to the underlying iterator
|
|
if hasattr(self._iterable, "take"):
|
|
self._iterable.take(n)
|
|
|
|
def __next__(self):
|
|
# Create consumer if not created yet
|
|
if self._consumer is None:
|
|
self._create_consumer()
|
|
|
|
# Notify the user if there is a data loading bottleneck
|
|
if self._queue.qsize() < min(2, max(1, self._queue.maxsize // 2)):
|
|
if time.time() - self.start_time > 5 * 60:
|
|
if (
|
|
self.warning_time is None
|
|
or time.time() - self.warning_time > 15 * 60
|
|
):
|
|
logger.debug(
|
|
"Data loading buffer is empty or nearly empty. This may "
|
|
"indicate a data loading bottleneck, and increasing the "
|
|
"number of workers (--num-workers) may help."
|
|
)
|
|
self.warning_time = time.time()
|
|
|
|
# Get next example
|
|
item = self._queue.get(True)
|
|
if isinstance(item, Exception):
|
|
raise item
|
|
if item is _sentinel:
|
|
raise StopIteration()
|
|
return item
|