def testBatchingSchemeMaxLength(self):
scheme = data_reader.batching_scheme(batch_size=20,
max_length=None,
min_length_bucket=8,
length_bucket_step=1.1,
drop_long_sequences=False)
self.assertGreater(scheme["max_length"], 10000)
scheme = data_reader.batching_scheme(batch_size=20,
max_length=None,
min_length_bucket=8,
length_bucket_step=1.1,
drop_long_sequences=True)
self.assertEqual(scheme["max_length"], 20)
scheme = data_reader.batching_scheme(batch_size=20,
max_length=15,
min_length_bucket=8,
length_bucket_step=1.1,
drop_long_sequences=True)
self.assertEqual(scheme["max_length"], 15)
scheme = data_reader.batching_scheme(batch_size=20,
max_length=15,
min_length_bucket=8,
length_bucket_step=1.1,
drop_long_sequences=False)
self.assertGreater(scheme["max_length"], 10000)
def testBatchingSchemeBuckets(self):
scheme = data_reader.batching_scheme(
batch_size=128,
max_length=0,
min_length_bucket=8,
length_bucket_step=1.1)
boundaries, batch_sizes = scheme["boundaries"], scheme["batch_sizes"]
self.assertEqual(len(boundaries), len(batch_sizes) - 1)
expected_boundaries = [
8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 24, 26, 28, 30,
33, 36, 39, 42, 46, 50, 55, 60, 66, 72, 79, 86, 94, 103, 113, 124
]
self.assertEqual(expected_boundaries, boundaries)
expected_batch_sizes = [
16, 12, 12, 8, 8, 8, 8, 8, 8, 6, 6, 6, 6, 4, 4, 4, 4, 4, 3, 3, 3, 3, 2,
2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1
]
self.assertEqual(expected_batch_sizes, batch_sizes)
scheme = data_reader.batching_scheme(
batch_size=128,
max_length=0,
min_length_bucket=8,
length_bucket_step=1.1,
shard_multiplier=2)
boundaries, batch_sizes = scheme["boundaries"], scheme["batch_sizes"]
self.assertAllEqual([bs * 2 for bs in expected_batch_sizes], batch_sizes)
self.assertEqual(expected_boundaries, boundaries)
scheme = data_reader.batching_scheme(
batch_size=128,
max_length=0,
min_length_bucket=8,
length_bucket_step=1.1,
length_multiplier=2)
boundaries, batch_sizes = scheme["boundaries"], scheme["batch_sizes"]
self.assertAllEqual([b * 2 for b in expected_boundaries], boundaries)
self.assertEqual([max(1, bs // 2)
for bs in expected_batch_sizes], batch_sizes)
def batch_fn(dataset,
training,
shapes,
target_names,
batch_size=32,
eval_batch_size=32,
bucket_batch_length=32,
bucket_max_length=256,
bucket_min_length=8,
bucket_length_step=1.1,
buckets=None):
"""Batching function."""
del target_names
# If bucketing is not specified, check if target shapes are variable.
cur_batch_size = batch_size if training else eval_batch_size
if buckets is None:
variable_target_shapes = False
target_shape = shapes[1]
for dim in target_shape:
if dim is None:
variable_target_shapes = True
tf.logging.info(
"Heuristically setting bucketing to %s based on shapes "
"of target tensors." % variable_target_shapes)
if variable_target_shapes:
batch_size_per_token = cur_batch_size * bucket_batch_length
scheme = data_reader.batching_scheme(batch_size_per_token,
bucket_max_length,
bucket_min_length,
bucket_length_step,
drop_long_sequences=training)
buckets = (scheme["boundaries"], scheme["batch_sizes"])
if buckets:
tf.logging.info("Bucketing with buckets %s." % str(buckets))
def example_length(_, target):
return tf.shape(target)[0]
boundaries, batch_sizes = buckets
dataset = dataset.apply(
tf.data.experimental.bucket_by_sequence_length(
example_length, boundaries, batch_sizes))
else:
dataset = dataset.padded_batch(cur_batch_size, shapes)
return dataset