class ReaderMockTest(unittest.TestCase):
def setUp(self):
self.reader = ReaderMock(TestSchema, schema_data_generator_example)
def test_simple_read(self):
"""Just a bunch of read and compares of all values to the expected values for their types and
shapes."""
# Read a bunch of entries from the dataset and compare the data to reference
for _ in range(10):
actual = dict(next(self.reader)._asdict())
for schema_field in TestSchema.fields.values():
if schema_field.numpy_dtype == Decimal:
self.assertTrue(
isinstance(actual[schema_field.name], Decimal))
else:
self.assertTrue(actual[schema_field.name].dtype.type is
schema_field.numpy_dtype)
self.assertEqual(len(actual[schema_field.name].shape),
len(schema_field.shape))
self.reader.stop()
self.reader.join()
@pytest.mark.forked
@create_tf_graph
def test_simple_read_tf(self):
"""Just a bunch of read and compares of all values to the expected values for their types
and shapes"""
reader_tensors = tf_tensors(self.reader)._asdict()
for schema_field in TestSchema.fields.values():
self.assertEqual(reader_tensors[schema_field.name].dtype,
_numpy_to_tf_dtypes(schema_field.numpy_dtype))
self.assertEqual(len(reader_tensors[schema_field.name].shape),
len(schema_field.shape))
# Read a bunch of entries from the dataset and compare the data to reference
with tf.Session() as sess:
for _ in range(10):
sess.run(reader_tensors)
self.reader.stop()
self.reader.join()
import numpy as np
import pytest
import six
from petastorm.predicates import in_lambda
from petastorm.reader import Reader
from petastorm.weighted_sampling_reader import WeightedSamplingReader
from petastorm.test_util.reader_mock import ReaderMock
from petastorm.unischema import Unischema, UnischemaField
from petastorm.workers_pool.dummy_pool import DummyPool
TestSchema = Unischema('TestSchema', [
UnischemaField('f1', np.int32, (), None, False),
])
reader0 = ReaderMock(TestSchema, lambda _: {'f1': 0})
reader1 = ReaderMock(TestSchema, lambda _: {'f1': 1})
reader2 = ReaderMock(TestSchema, lambda _: {'f1': 2})
def _count_mixed(readers, probabilities, num_of_reads):
result = len(probabilities) * [0]
with WeightedSamplingReader(readers, probabilities) as mixer:
for _ in six.moves.xrange(num_of_reads):
reader_index = next(mixer).f1
result[reader_index] += 1
return result
def setUp(self):
self.reader = ReaderMock(TestSchema, schema_data_generator_example)