def test_noop_shuffling_buffer():
"""Noop should not do any shuffling. Add/retrieve some items while checking correcness of can_retrieve"""
q = NoopShufflingBuffer()
# Empty buffer. Can add, can not retrieve with zero size
assert q.size == 0
assert q.can_add()
assert not q.can_retrieve()
# Try adding some items. Check queue size and can_retrieve indicator
q.add_many([1])
assert q.can_add()
assert q.size == 1
q.add_many([2, 3])
assert q.size == 3
assert 1 == q.retrieve()
assert q.can_retrieve()
assert 2 == q.retrieve()
assert 3 == q.retrieve()
assert not q.can_retrieve()
# No effect is expected in noop implementation
q.finish()
class DataLoader(object):
"""
A data loader adaptor for ``torch.utils.data.DataLoader``.
This class iterates and returns items from the Reader in batches.
This loader can be used as an iterator and will terminate when the reader used in the construction of the class
runs out of samples.
"""
def __init__(self, reader, batch_size=1, collate_fn=decimal_friendly_collate,
shuffling_queue_capacity=0):
"""
Initializes a data loader object, with a default collate.
Number of epochs is defined by the configuration of the reader argument.
An optional shuffling queue is created if shuffling_queue_capacity is greater than 0. No samples will be
returned to a user by the ``DataLoader`` until the queue is full. After that, batches of `batch_size`
will be created by uniformly sampling the shuffling queue. Once no more samples are available from the data
reader, the shuffling queue is allowed to be consumed till no further samples are available.
Note that the last returned batch could have less then ``batch_size`` samples.
NOTE: if you are using ``make_batch_reader``, this shuffling queue will be randomizing the order of the
entire batches and not changing the order of elements within a batch. This is likely not what you intend to do.
:param reader: petastorm Reader instance
:param batch_size: the number of items to return per batch; factored into the len() of this reader
:param collate_fn: an optional callable to merge a list of samples to form a mini-batch.
:param shuffling_queue_capacity: Queue capacity is passed to the underlying :class:`tf.RandomShuffleQueue`
instance. If set to 0, no suffling will be done.
"""
self.reader = reader
self.batch_size = batch_size
self.collate_fn = collate_fn
# _batch_acc accumulates samples for a single batch.
self._batch_acc = []
if shuffling_queue_capacity > 0:
# We can not know what is the reasonable number to use for the extra capacity, so we set a huge number
# and give up on the unbound growth protection mechanism.
min_after_dequeue = shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=100000000)
else:
self._shuffling_buffer = NoopShufflingBuffer()
def __iter__(self):
"""
The Data Loader iterator stops the for-loop when reader runs out of samples.
"""
# As we iterate over incoming samples, we are going to store them in `self._batch_acc`, until we have a batch of
# the requested batch_size ready.
keys = None
for row in self.reader:
# Default collate does not work nicely on namedtuples and treat them as lists
# Using dict will result in the yielded structures being dicts as well
row_as_dict = row._asdict()
keys = row_as_dict.keys()
# Promote some types that are incompatible with pytorch to be pytorch friendly.
_sanitize_pytorch_types(row_as_dict)
# Add rows to shuffling buffer
if not self.reader.is_batched_reader:
self._shuffling_buffer.add_many([row_as_dict])
else:
# Transposition:
# row_as_dict: {'a': [1,2,3], 'b':[4,5,6]}
# row_group_as_tuple: [(1, 4), (2, 5), (3, 6)]
# The order within a tuple is defined by key order in 'keys'
row_group_as_tuple = list(zip(*(row_as_dict[k] for k in keys)))
# Adding data as 'row-by-row' into a shuffling buffer. This is a pretty
# slow implementation though. Probably can comeup with a faster way to shuffle,
# perhaps at the expense of a larger memory consumption...
self._shuffling_buffer.add_many(row_group_as_tuple)
# _yield_batches will emit as much batches as are allowed by the shuffling_buffer (RandomShufflingBuffer
# will avoid underflowing below a certain number of samples to guarantee some samples decorrelation)
for batch in self._yield_batches(keys):
yield batch
# Once reader can not read new rows, we might still have a bunch of rows waiting in the shuffling buffer.
# Telling shuffling buffer that we are finished allows to deplete the buffer completely, regardless its
# min_after_dequeue setting.
self._shuffling_buffer.finish()
for batch in self._yield_batches(keys):
yield batch
# Yield the last and partial batch
if self._batch_acc:
yield self.collate_fn(self._batch_acc)
def _yield_batches(self, keys):
while self._shuffling_buffer.can_retrieve():
post_shuffled_row = self._shuffling_buffer.retrieve()
if not isinstance(post_shuffled_row, dict):
# This is for the case of batched reads. Here we restore back the
# dictionary format of records
post_shuffled_row = dict(zip(keys, post_shuffled_row))
self._batch_acc.append(post_shuffled_row)
# Batch is ready? Collate and emmit
if len(self._batch_acc) == self.batch_size:
yield self.collate_fn(self._batch_acc)
self._batch_acc = []
# Functions needed to treat data loader as a context manager
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.reader.stop()
self.reader.join()
class DataLoader(object):
"""
A data loader adaptor for ``torch.utils.data.DataLoader``.
This class iterates and returns items from the Reader in batches.
This loader can be used as an iterator and will terminate when the reader used in the construction of the class
runs out of samples.
"""
def __init__(self,
reader,
batch_size=1,
collate_fn=decimal_friendly_collate,
shuffling_queue_capacity=0):
"""
Initializes a data loader object, with a default collate.
Number of epochs is defined by the configuration of the reader argument.
An optional shuffling queue is created if shuffling_queue_capacity is greater than 0. No samples will be
returned to a user by the ``DataLoader`` until the queue is full. After that, batches of `batch_size`
will be created by uniformly sampling the shuffling queue. Once no more samples are available from the data
reader, the shuffling queue is allowed to be consumed till no further samples are available.
Note that the last returned batch could have less then ``batch_size`` samples.
NOTE: if you are using ``make_batch_reader``, this shuffling queue will be randomizing the order of the
entire batches and not changing the order of elements within a batch. This is likely not what you intend to do.
:param reader: petastorm Reader instance
:param batch_size: the number of items to return per batch; factored into the len() of this reader
:param collate_fn: an optional callable to merge a list of samples to form a mini-batch.
:param shuffling_queue_capacity: Queue capacity is passed to the underlying :class:`tf.RandomShuffleQueue`
instance. If set to 0, no suffling will be done.
"""
self.reader = reader
self.batch_size = batch_size
self.collate_fn = collate_fn
if shuffling_queue_capacity > 0:
# We always add and remove one sample from the shuffling buffer in the steady state. min_after_dequeue
# is important when enqueue and dequeue rates are not the same.
# By setting `min_after_dequeue = shuffling_queue_capacity - 1` we will just keep the buffer full.
min_after_dequeue = shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(
shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=0)
else:
self._shuffling_buffer = NoopShufflingBuffer()
def __iter__(self):
"""
The Data Loader iterator stops the for-loop when reader runs out of samples.
"""
batch = []
for row in self.reader:
# Default collate does not work nicely on namedtuples and treat them as lists
# Using dict will result in the yielded structures being dicts as well
row_as_dict = row._asdict()
# Promote some types that are incompatible with pytorch to be pytorch friendly.
_sanitize_pytorch_types(row_as_dict)
# Add rows to shuffling buffer
self._shuffling_buffer.add_many([row_as_dict])
# If the shuffling buffer is ready, pull the resampled rows out and try add to our accumulated batch
while self._shuffling_buffer.can_retrieve():
post_shuffled_row = self._shuffling_buffer.retrieve()
batch.append(post_shuffled_row)
# Batch is ready? Collate and emmit
if len(batch) == self.batch_size:
yield self.collate_fn(batch)
batch = []
# Once reader can not emit new rows, we might still have a bunch of rows waiting in the shuffling buffer.
# Telling shuffling buffer thatt we are finished allows to deplete the buffer completely, regardless its
# min_after_dequeue setting.
self._shuffling_buffer.finish()
# Some code duplication with the main reading loop. Not sure how to reorganize to avoid this without
# overcomplicating the code too much.
while self._shuffling_buffer.can_retrieve():
post_shuffled_row = self._shuffling_buffer.retrieve()
batch.append(post_shuffled_row)
# Batch is ready? Collate and emmit
if len(batch) == self.batch_size:
yield self.collate_fn(batch)
batch = []
if batch:
yield self.collate_fn(batch)
# Functions needed to treat data loader as a context manager
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.reader.stop()
self.reader.join()
class DataLoader(object):
"""
A data loader adaptor for ``torch.utils.data.DataLoader``.
This class iterates and returns items from the Reader in batches.
This loader can be used as an iterator and will terminate when the reader used in the construction of the class
runs out of samples.
"""
def __init__(self,
reader,
batch_size=1,
collate_fn=torch.as_tensor,
shuffling_queue_capacity=0):
"""
Initializes a data loader object, with a default collate.
Number of epochs is defined by the configuration of the reader argument.
An optional shuffling queue is created if shuffling_queue_capacity is greater than 0. No samples will be
returned to a user by the ``DataLoader`` until the queue is full. After that, batches of `batch_size`
will be created by uniformly sampling the shuffling queue. Once no more samples are available from the data
reader, the shuffling queue is allowed to be consumed till no further samples are available.
Note that the last returned batch could have less then ``batch_size`` samples.
NOTE: if you are using ``make_batch_reader``, this shuffling queue will be randomizing the order of the
entire batches and not changing the order of elements within a batch. This is likely not what you intend to do.
:param reader: petastorm Reader instance
:param batch_size: the number of items to return per batch; factored into the len() of this reader
:param collate_fn: an optional callable to merge a list of samples to form a mini-batch.
:param shuffling_queue_capacity: Queue capacity is passed to the underlying :class:`tf.RandomShuffleQueue`
instance. If set to 0, no suffling will be done.
"""
self.reader = reader
self.batch_size = batch_size
self.collate_fn = collate_fn
# _batch_acc accumulates samples for a single batch.
self._batch_acc = []
if shuffling_queue_capacity > 0:
# We can not know what is the reasonable number to use for the extra capacity, so we set a huge number
# and give up on the unbound growth protection mechanism.
min_after_dequeue = shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(
shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=100000000)
else:
self._shuffling_buffer = NoopShufflingBuffer()
def __iter__(self):
"""
The Data Loader iterator stops the for-loop when reader runs out of samples.
"""
# As we iterate over incoming samples, we are going to store them in `self._batch_acc`, until we have a batch of
# the requested batch_size ready.
keys = None
for row in self.reader:
# Default collate does not work nicely on namedtuples and treat them as lists
# Using dict will result in the yielded structures being dicts as well
row_as_dict = row._asdict()
keys = row_as_dict.keys()
# Promote some types that are incompatible with pytorch to be pytorch friendly.
_sanitize_pytorch_types(row_as_dict)
# Add rows to shuffling buffer
for k, v in row_as_dict.items():
if not self.reader.is_batched_reader:
row_as_dict[k] = self.collate_fn([v])
else:
row_as_dict[k] = self.collate_fn(v)
self._shuffling_buffer.add_many(row_as_dict)
# _yield_batches will emit as much batches as are allowed by the shuffling_buffer (RandomShufflingBuffer
# will avoid underflowing below a certain number of samples to guarantee some samples decorrelation)
for batch in self._yield_batches(keys):
yield batch
# Once reader can not read new rows, we might still have a bunch of rows waiting in the shuffling buffer.
# Telling shuffling buffer that we are finished allows to deplete the buffer completely, regardless its
# min_after_dequeue setting.
self._shuffling_buffer.finish()
for batch in self._yield_batches(keys):
yield batch
def _yield_batches(self, keys):
while self._shuffling_buffer.can_retrieve(self.batch_size):
batch = self._shuffling_buffer.retrieve(self.batch_size)
yield batch
# Functions needed to treat data loader as a context manager
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.reader.stop()
self.reader.join()
