def reduce_func(unused_key, window: tf.data.Dataset):
return window.padded_batch(
batch_size,
padded_shapes=(SourceData(
key=tf.TensorShape([]),
waveform=tf.TensorShape([None, 1]),
waveform_length=tf.TensorShape([]),
speaker_id=tf.TensorShape([]),
age=tf.TensorShape([]),
gender=tf.TensorShape([]),
),
SourceData(
key=tf.TensorShape([]),
waveform=tf.TensorShape([None, 1]),
waveform_length=tf.TensorShape([]),
speaker_id=tf.TensorShape([]),
age=tf.TensorShape([]),
gender=tf.TensorShape([]),
)),
padding_values=(SourceData(
key="",
waveform=tf.to_float(0),
waveform_length=tf.to_int64(0),
speaker_id=tf.to_int64(0),
age=tf.to_int64(0),
gender=tf.to_int64(0),
),
SourceData(
key="",
waveform=tf.to_float(0),
waveform_length=tf.to_int64(0),
speaker_id=tf.to_int64(0),
age=tf.to_int64(0),
gender=tf.to_int64(0),
)))
def reduce_func(unused_key, window: tf.data.Dataset):
return window.padded_batch(
batch_size,
padded_shapes=(SourceData(
id=tf.TensorShape([]),
key=tf.TensorShape([]),
mel=tf.TensorShape([None, self.hparams.num_mels]),
mel_length=tf.TensorShape([]),
mel_width=tf.TensorShape([]),
text=tf.TensorShape([]),
),
TargetData(
id=tf.TensorShape([]),
key=tf.TensorShape([]),
waveform=tf.TensorShape([None, 1]),
waveform_length=tf.TensorShape([]),
)),
padding_values=(SourceData(
id=tf.to_int64(0),
key="",
mel=tf.to_float(0),
mel_length=tf.to_int64(0),
mel_width=tf.to_int64(0),
text="",
),
TargetData(
id=tf.to_int64(0),
key="",
waveform=tf.to_float(0),
waveform_length=tf.to_int64(0),
)))
def process(self, dataset: tf.data.Dataset, batch_size: int):
dataset = dataset.map(self.parse, num_parallel_calls=AUTOTUNE)
if self.cache:
dataset = dataset.cache()
if self.shuffle:
dataset = dataset.shuffle(self.buffer_size,
reshuffle_each_iteration=True)
# PADDED BATCH the dataset
dataset = dataset.padded_batch(
batch_size=batch_size,
padded_shapes=(
tf.TensorShape([]),
tf.TensorShape(self.speech_featurizer.shape),
tf.TensorShape([]),
tf.TensorShape([None]),
tf.TensorShape([]),
tf.TensorShape([None]),
tf.TensorShape([]),
),
padding_values=("", 0., 0, self.text_featurizer.blank, 0,
self.text_featurizer.blank, 0),
drop_remainder=self.drop_remainder)
# PREFETCH to improve speed of input length
dataset = dataset.prefetch(AUTOTUNE)
self.total_steps = get_num_batches(self.total_steps, batch_size)
return dataset
def batch_and_split(dataset: tf.data.Dataset, max_sequence_length: int,
batch_size: int) -> tf.data.Dataset:
return dataset.padded_batch(
batch_size,
padded_shapes=[max_sequence_length + 1
]).map(split_input_target,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
def get_bucket_iter(self,
dataset: tf.data.Dataset,
batch_size=32,
train=True) -> tf.data.Dataset:
padded_shapes = self._padded_shapes()
padding_values = self._padding_values()
if train:
bucket_boundaries = self._bucket_boundaries(batch_size)
bucket_batch_sizes = [batch_size] * (len(bucket_boundaries) + 1)
dataset = dataset.apply(
tf.data.experimental.bucket_by_sequence_length(
self.element_length_func,
bucket_boundaries,
bucket_batch_sizes,
padded_shapes=padded_shapes,
padding_values=padding_values))
dataset = dataset.shuffle(100)
else:
dataset = dataset.padded_batch(batch_size,
padded_shapes=padded_shapes)
dataset = dataset.map(self._collate_fn,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
return dataset
def batch_dataset(dataset: tf.data.Dataset, feature_type2name,
feature_name2num, batch_size):
"""Performs batching on ranking dataset
When there's no sparse features, padded_batch() is enough. When there are sparse features, we use batching function specific for sparse features
"""
padded_shapes, padded_values = _get_padded_shapes_and_values(
feature_type2name, feature_name2num)
# Use padded_batch() if no sparse features
if InputFtrType.SPARSE_FTRS_COLUMN_NAMES not in feature_type2name and InputFtrType.SHALLOW_TOWER_SPARSE_FTRS_COLUMN_NAMES not in feature_type2name:
# drop_remainder=True to avoid input batch_size=0 issue in evaluation mode in multi gpu training
return dataset.padded_batch(batch_size,
padded_shapes=padded_shapes,
padding_values=padded_values,
drop_remainder=True)
sparse_batch_fn = partial(_sparse_batch_fn,
feature_type2name=feature_type2name,
padded_shapes=padded_shapes,
padded_values=padded_values,
batch_size=batch_size)
# drop_remainder=True to avoid input batch_size=0 issue in evaluation mode in multi gpu training
return dataset.window(batch_size,
drop_remainder=True).flat_map(sparse_batch_fn)
def batch_dataset(dataset: tf.data.Dataset,
model: LineRecognizer,
batch_size=32,
bucket_boundaries=None,
padded=True):
# add image widths and text length
dataset = dataset.map(lambda i, t: (i, tf.shape(i)[
1], t, tf.strings.length(t, unit='UTF8_CHAR')))
dataset = dataset.map(lambda image, width, text, length:
(image, width, model.encoder.encode(text), length))
output_shapes = (model.image_shape, [], [None], [])
if bucket_boundaries:
if isinstance(batch_size, int):
batch_size = [batch_size] * (len(bucket_boundaries) + 1)
dataset = dataset.apply(
tf.data.experimental.bucket_by_sequence_length(
lambda i, w, label, length: w,
bucket_boundaries=bucket_boundaries,
bucket_batch_sizes=batch_size,
padded_shapes=output_shapes))
elif padded:
dataset = dataset.padded_batch(batch_size=batch_size,
padded_shapes=output_shapes)
else:
dataset = dataset.batch(batch_size)
return dataset
def run_evaluate(self, data: tf.data.Dataset):
data = data.padded_batch(self.args.batch)
metrics = defaultdict(list)
for batch_id, batch in enumerate(data):
out = self(batch, training=True)
preds = self.extract_preds(out, batch)
for k, v in self.compute_metrics(preds, batch).items():
metrics[k].append(v)
return {k: sum(v) / len(v) for k, v in metrics.items()}
def apply(self, dataset: tf.data.Dataset, mode: str = None):
# pylint: disable=unused-argument
padded_shapes = {field.name: field.shape for field in self.fields}
padding_values = {
field.name: tf.constant(field.default, field.dtype)
for field in self.fields
}
return dataset.padded_batch(
self.batch_size,
padded_shapes=padded_shapes,
padding_values=padding_values,
drop_remainder=self.drop_remainder,
)
def _hook_dataset_post_precessing(self, ds: tf.data.Dataset,
batch_size: int):
ds = ds.padded_batch(batch_size=batch_size, padded_shapes=[None])
ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
def add_padding_mask(source):
enc_padding_mask = create_padding_mask_fm(source)
return (source, enc_padding_mask), source
ds = ds.map(map_func=add_padding_mask) # Add pad mask
# For masked language model task, all datasets are masked
return self._apply_mask_for_mlm(ds=ds, vocab_size=self._vocab_size)
def pad_dataset(data: tf.data.Dataset, bs: int, max_length: int, train: bool):
data = data.padded_batch(batch_size=bs,
padded_shapes=({
"input_ids": [max_length],
"segment_ids": [max_length],
"input_mask": [max_length]
}, []),
padding_values=({
"input_ids": 0,
"segment_ids": 0,
"input_mask": 0,
}, 0),
drop_remainder=True)
return data
def get_data_iter(self,
dataset: tf.data.Dataset,
batch_size=32,
train=True) -> tf.data.Dataset:
padded_shapes = self._padded_shapes()
padding_values = self._padding_values()
if train:
dataset = dataset.shuffle(batch_size * 100)
dataset = dataset.padded_batch(batch_size,
padded_shapes=padded_shapes,
padding_values=padding_values)
dataset = dataset.map(self._collate_fn,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
return dataset
def prepare_dataset(
dataset: tf.data.Dataset,
model_image_size: Tuple[int, int],
augmentation_fn: Optional[ImageDataMapFn] = None,
num_epochs: Optional[int] = None,
batch_size: Optional[int] = None,
shuffle_buffer_size: Optional[int] = None,
num_parallel_calls: Optional[int] = None,
prefetch_buffer_size: Optional[int] = None,
prefetch_to_device: Optional[str] = None,
) -> tf.data.Dataset:
# apply data augmentation:
if augmentation_fn is not None:
dataset = dataset.map(
map_image_data(augmentation_fn),
num_parallel_calls=num_parallel_calls,
)
if shuffle_buffer_size is not None:
dataset = dataset.shuffle(buffer_size=shuffle_buffer_size)
dataset = dataset.repeat(num_epochs)
dataset = dataset.map(
map_image_data(prepare_for_batching(model_image_size)),
num_parallel_calls=num_parallel_calls,
)
# batching and padding
if batch_size is not None:
dataset = dataset.padded_batch(
batch_size=batch_size,
padded_shapes=get_padding_shapes(
dataset, spatial_image_shape=model_image_size),
drop_remainder=True,
)
# try to prefetch dataset on certain device
if prefetch_to_device is not None:
prefetch_fn = tf.data.experimental.prefetch_to_device(
device=prefetch_to_device, buffer_size=prefetch_buffer_size)
dataset = dataset.apply(prefetch_fn)
else:
if prefetch_buffer_size is not None:
dataset = dataset.prefetch(buffer_size=prefetch_buffer_size)
return dataset
def preprocess_wmt_data(dataset: tf.data.Dataset,
shuffle: bool,
num_epochs: Optional[int] = 1,
pack_examples: bool = True,
shuffle_buffer_size: int = 1024,
max_length: int = 512,
batch_size: int = 256,
drop_remainder: bool = True,
prefetch_size: int = AUTOTUNE):
"""Shuffle and batch/pack the given dataset."""
def length_filter(max_len):
def filter_fn(x):
source, target = x['inputs'], x['targets']
l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0])
return tf.less(l, max_len + 1)
return filter_fn
if max_length > 0:
dataset = dataset.filter(length_filter(max_length))
if shuffle:
dataset = dataset.shuffle(shuffle_buffer_size)
dataset = dataset.repeat(num_epochs)
if pack_examples:
dataset = pack_dataset(dataset, max_length)
dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)
else: # simple (static-shape) padded batching
dataset = dataset.padded_batch(batch_size,
padded_shapes={
'inputs': max_length,
'targets': max_length
},
padding_values={
'inputs': 0,
'targets': 0
},
drop_remainder=drop_remainder)
if prefetch_size:
dataset = dataset.prefetch(prefetch_size)
return dataset
def reduce_func(unused_key, window: tf.data.Dataset):
return window.padded_batch(batch_size, padded_shapes=(
SourceData(
id=tf.TensorShape([]),
key=tf.TensorShape([]),
source=tf.TensorShape([None]),
source_length=tf.TensorShape([]),
speaker_id=tf.TensorShape([]),
age=tf.TensorShape([]),
gender=tf.TensorShape([]),
text=tf.TensorShape([]),
),
MelData(
id=tf.TensorShape([]),
key=tf.TensorShape([]),
mel=tf.TensorShape([None, self.hparams.num_mels]),
mel_width=tf.TensorShape([]),
target_length=tf.TensorShape([]),
done=tf.TensorShape([None]),
spec_loss_mask=tf.TensorShape([None]),
binary_loss_mask=tf.TensorShape([None]),
)), padding_values=(
SourceData(
id=tf.to_int64(0),
key="",
source=tf.to_int64(0),
source_length=tf.to_int64(0),
speaker_id=tf.to_int64(0),
age=tf.to_int64(0),
gender=tf.to_int64(-1),
text="",
),
MelData(
id=tf.to_int64(0),
key="",
mel=tf.to_float(self.hparams.silence_mel_level_db),
mel_width=tf.to_int64(0),
target_length=tf.to_int64(0),
done=tf.to_float(1),
spec_loss_mask=tf.to_float(0),
binary_loss_mask=tf.to_float(0),
)))
def combine_samples_to_batch(self, data: tf.data.Dataset,
batch_size: int) -> tf.data.Dataset:
"""
Apply batching on data
Parameters
----------
data
data
batch_size
batch size
Returns
-------
data_batch
batched data
"""
data = data.padded_batch(batch_size,
data.output_shapes,
drop_remainder=self.fix_batch_dimension)
return data
def preprocess_wmt_data(dataset: tf.data.Dataset,
data_rng,
train: bool,
shuffle_buffer_size: int = 1024,
max_length: int = 512,
per_device_batch_size: int = 256):
"""Shuffle and batch/pack the given dataset."""
def length_filter(max_len):
def filter_fn(x):
source, target = x['inputs'], x['targets']
l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0])
return tf.less(l, max_len + 1)
return filter_fn
if max_length > 0:
dataset = dataset.filter(length_filter(max_length))
if train:
dataset = dataset.shuffle(shuffle_buffer_size, seed=data_rng[0])
dataset = dataset.repeat()
dataset = pack_dataset(dataset, max_length)
dataset = dataset.batch(per_device_batch_size, drop_remainder=train)
else: # simple (static-shape) padded batching
dataset = dataset.padded_batch(per_device_batch_size,
padded_shapes={
'inputs': max_length,
'targets': max_length
},
padding_values={
'inputs': 0,
'targets': 0
},
drop_remainder=train)
dataset = dataset.prefetch(AUTOTUNE)
return dataset
def pack_dataset(dataset: tf.data.Dataset,
key2length: Union[int, Dict[str, int]],
keys: Optional[List[str]] = None) -> tf.data.Dataset:
"""Creates a 'packed' version of a dataset on-the-fly.
Adapted from the mesh-tf implementation.
This is meant to replace the irritation of having to create a separate
"packed" version of a dataset to train efficiently on TPU.
Each example in the output dataset represents several examples in the
input dataset.
For each key in the input dataset, two additional keys are created:
<key>_segmentation: an int32 tensor identifying the parts
representing the original example.
<key>_position: an int32 tensor identifying the position within the original
example.
Example:
Two input examples get combined to form an output example.
The input examples are:
{"inputs": [8, 7, 1, 0], "targets":[4, 1, 0]}
{"inputs": [2, 3, 4, 1], "targets":[5, 6, 1]}
The output example is:
{
"inputs": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]
"inputs_segmentations": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]
"inputs_positions": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]
"targets": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]
"targets_segmentations": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]
"targets_positions": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]
}
0 represents padding in both the inputs and the outputs.
Sequences in the incoming examples are truncated to length "length", and the
sequences in the output examples all have fixed (padded) length "length".
Args:
dataset: a tf.data.Dataset
key2length: an integer, or a dict from feature-key to integer
keys: a list of strings (e.g. ["inputs", "targets"])
Returns:
a tf.data.Dataset
"""
shapes = tf.nest.map_structure(lambda spec: spec.shape,
dataset.element_spec)
if keys is None:
keys = list(shapes.keys())
for k in keys:
if k not in shapes:
raise ValueError(
f'Key {k} not found in dataset. Available keys are {shapes.keys()}'
)
if not shapes[k].is_compatible_with(tf.TensorShape([None])):
raise ValueError('Tensors to be packed must be one-dimensional.')
# make sure that the length dictionary contains all keys as well as the
# keys suffixed by "_segmentation" and "_position"
if isinstance(key2length, int):
key2length = {k: key2length for k in keys}
for k in keys:
for suffix in ['_segmentation', '_position']:
key2length[k + suffix] = key2length[k]
# trim to length
dataset = dataset.map(lambda x: {k: x[k][:key2length[k]]
for k in keys},
num_parallel_calls=AUTOTUNE)
# Setting batch_size=length ensures that the concatenated sequences (if they
# have length >=1) are sufficient to fill at least one packed example.
batch_size = max(key2length.values())
dataset = dataset.padded_batch(batch_size,
padded_shapes={k: [-1]
for k in keys})
dataset = _pack_with_tf_ops(dataset, keys, key2length)
# Set the Tensor shapes correctly since they get lost in the process.
def my_fn(x):
return {k: tf.reshape(v, [key2length[k]]) for k, v in x.items()}
return dataset.map(my_fn, num_parallel_calls=AUTOTUNE)
def dataset_batch(ds_single: tf.data.Dataset, batch_shape, batch_size=5):
ds = ds_single.padded_batch(batch_size=batch_size,
padded_shapes=batch_shape,
drop_remainder=False)
return ds
def __batch_and_prefetch(self, dataset: tf.data.Dataset) -> tf.data.Dataset:
return dataset.padded_batch(
self.batch_size,
padded_shapes=([None, None, 3],
[self.no_classes, ])).prefetch(
buffer_size=self.buffer_size)
