def test_batch(self):
""" test batched dataset
"""
batchsize = 2
mapper = tf.build_mapper(self.ops)
ds = build_source(self.sc_config)
mapped_ds = tf.map(ds, mapper)
batched_ds = tf.batch(mapped_ds, batchsize, True)
for sample in batched_ds:
out = sample
self.assertEqual(len(out), batchsize)
def test_map(self):
""" test transformer.map
"""
mapper = tf.build_mapper(self.ops)
ds = build_source(self.sc_config)
mapped_ds = tf.map(ds, mapper)
ct = 0
for sample in mapped_ds:
self.assertTrue(type(sample[0]) is np.ndarray)
ct += 1
self.assertEqual(ct, mapped_ds.size())
def test_transform_with_abnormal_worker(self):
""" test dataset transform with abnormally exit process
"""
samples = list(range(1000))
ds = MemorySource(samples)
def _mapper(sample):
if sample == 3:
sys.exit(1)
return 2 * sample
worker_conf = {'WORKER_NUM': 2, 'use_process': True}
mapped = tf.map(ds, _mapper, worker_conf)
ct = 0
for i, d in enumerate(mapped):
ct += 1
self.assertTrue(d / 2 in samples)
self.assertEqual(len(samples) - 1, ct)
def test_transform_with_delay_worker(self):
""" test dataset transform with delayed process
"""
samples = list(range(1000))
ds = MemorySource(samples)
def _mapper(sample):
if sample == 3:
time.sleep(30)
return 2 * sample
worker_conf = {'WORKER_NUM': 2, 'use_process': True}
mapped = tf.map(ds, _mapper, worker_conf)
ct = 0
for i, d in enumerate(mapped):
ct += 1
self.assertTrue(d / 2 in samples)
self.assertEqual(len(samples), ct)
def test_parallel_map(self):
""" test transformer.map with concurrent workers
"""
mapper = tf.build_mapper(self.ops)
ds = build_source(self.sc_config)
worker_conf = {'WORKER_NUM': 2, 'use_process': True}
mapped_ds = tf.map(ds, mapper, worker_conf)
ct = 0
for sample in mapped_ds:
self.assertTrue(type(sample[0]) is np.ndarray)
ct += 1
self.assertTrue(mapped_ds.drained())
self.assertEqual(ct, mapped_ds.size())
mapped_ds.reset()
ct = 0
for sample in mapped_ds:
self.assertTrue(type(sample[0]) is np.ndarray)
ct += 1
self.assertEqual(ct, mapped_ds.size())