def test_basic_two_workers(self):
data = self.get_test_data(1000)
batch_size = 100
num_shards = 2
reader = NpArrayReader(data,
batch_size=batch_size,
num_shards=num_shards)
splits = [
reader._get_split(reader.data, i, batch_size) for i in range(10)
]
streamer = DataStreamer(reader, num_workers=num_shards)
for _i, batch in enumerate(streamer):
match = False
for split in splits:
try:
self.assert_batch_equal(split, batch, 0, batch_size)
except Exception:
pass
else:
match = True
break
self.assertTrue(match)
self.assertEqual(9, _i)
def test_shard_drop_small(self):
n = 1000
data = self.get_test_data(n)
batch_size = 100
num_shards = 3
shard_size = n // num_shards + 1
reader = NpArrayReader(data, batch_size=batch_size, num_shards=num_shards)
num_batches = 0
for shard in range(num_shards):
for i, batch in enumerate(reader.get_shard(shard)):
self.assert_batch_equal(
data, batch, shard * shard_size + i * batch_size, batch_size
)
num_batches += 1
self.assertEqual(9, num_batches)
def test_shard_drop_small(self):
n = 1000
data = self.get_test_data(n)
batch_size = 100
num_shards = 3
shard_size = n // num_shards + 1
reader = NpArrayReader(data, batch_size=batch_size, num_shards=num_shards)
num_batches = 0
for shard in range(num_shards):
for i, batch in enumerate(reader.get_shard(shard)):
self.assert_batch_equal(
data, batch, shard * shard_size + i * batch_size, batch_size
)
num_batches += 1
self.assertEqual(9, num_batches)
def test_drop_small(self):
data = self.get_test_data(999)
batch_size = 100
reader = NpArrayReader(data, batch_size=batch_size)
for i, batch in enumerate(reader):
self.assert_batch_equal(data, batch, i * batch_size, batch_size)
self.assertEqual(8, i)
def test_basic(self):
data = self.get_test_data(1000)
batch_size = 100
reader = NpArrayReader(data, batch_size=batch_size)
for i, batch in enumerate(reader):
self.assert_batch_equal(data, batch, i * batch_size, batch_size)
self.assertEqual(9, i)
def test_drop_small_one_worker(self):
data = self.get_test_data(999)
batch_size = 100
reader = NpArrayReader(data, batch_size=batch_size, num_shards=1)
streamer = DataStreamer(reader, num_workers=1)
for i, batch in enumerate(streamer):
self.assert_batch_equal(data, batch, i * batch_size, batch_size)
self.assertEqual(8, i)
def test_not_drop_small(self):
data = self.get_test_data(999)
batch_size = 100
reader = NpArrayReader(data, batch_size=batch_size, drop_small=False)
streamer = DataStreamer(reader)
for i, batch in enumerate(streamer):
self.assert_batch_equal(data, batch, i * batch_size,
batch_size if i != 9 else 99)
self.assertEqual(9, i)
def test_basic_two_workers(self):
data = self.get_test_data(1000)
batch_size = 100
num_shards = 2
reader = NpArrayReader(data, batch_size=batch_size, num_shards=num_shards)
splits = [reader._get_split(reader.data, i, batch_size) for i in range(10)]
streamer = DataStreamer(reader, num_workers=num_shards)
for _i, batch in enumerate(streamer):
match = False
for split in splits:
try:
self.assert_batch_equal(split, batch, 0, batch_size)
except Exception:
pass
else:
match = True
break
self.assertTrue(match)
self.assertEqual(9, _i)