def _testFilterByLength(
self,
features_length,
labels_length,
maximum_features_length=None,
maximum_labels_length=None,
filtered=True,
):
dataset = tf.data.Dataset.zip(
(
tf.data.Dataset.from_tensors(tf.constant(features_length)),
tf.data.Dataset.from_tensors(tf.constant(labels_length)),
)
)
dataset = dataset.apply(
dataset_util.filter_examples_by_length(
maximum_features_length=maximum_features_length,
maximum_labels_length=maximum_labels_length,
features_length_fn=lambda _: features_length,
labels_length_fn=lambda _: labels_length,
)
)
iterator = iter(dataset)
if filtered:
with self.assertRaises(StopIteration):
next(iterator)
else:
next(iterator)
def make_training_dataset(self,
features_file,
labels_file,
batch_size,
batch_type="examples",
batch_multiplier=1,
batch_size_multiple=1,
shuffle_buffer_size=None,
length_bucket_width=None,
maximum_features_length=None,
maximum_labels_length=None,
single_pass=False,
num_shards=1,
shard_index=0,
num_threads=4,
prefetch_buffer_size=None,
cardinality_multiple=1,
weights=None):
"""See :meth:`opennmt.inputters.ExampleInputter.make_training_dataset`."""
_ = labels_file
dataset = self.make_dataset(features_file, training=True)
if weights is not None:
dataset = (dataset, weights)
transform_fns = []
map_func = lambda x: self.make_features(element=x, training=True)
transform_fns.append(lambda dataset:
dataset.map(map_func,
num_parallel_calls=num_threads or 4))
transform_fns.append(dataset_util.filter_examples_by_length(
maximum_features_length=maximum_features_length,
maximum_labels_length=maximum_labels_length,
features_length_fn=self.get_length))
dataset = dataset_util.training_pipeline(
batch_size,
batch_type=batch_type,
batch_multiplier=batch_multiplier,
batch_size_multiple=batch_size_multiple,
transform_fns=transform_fns,
length_bucket_width=length_bucket_width,
single_pass=single_pass,
num_shards=num_shards,
shard_index=shard_index,
num_threads=num_threads,
shuffle_buffer_size=shuffle_buffer_size,
prefetch_buffer_size=prefetch_buffer_size,
cardinality_multiple=cardinality_multiple)(dataset)
return dataset
def make_training_dataset(self,
features_file,
labels_file,
batch_size,
batch_type="examples",
batch_multiplier=1,
batch_size_multiple=1,
shuffle_buffer_size=None,
length_bucket_width=None,
maximum_features_length=None,
maximum_labels_length=None,
single_pass=False,
num_shards=1,
shard_index=0,
num_threads=4,
prefetch_buffer_size=None,
cardinality_multiple=1,
weights=None):
"""Builds a dataset to be used for training. It supports the full training
pipeline, including:
* sharding
* shuffling
* filtering
* bucketing
* prefetching
Args:
features_file: The source file or a list of training source files.
labels_file: The target file or a list of training target files.
batch_size: The batch size to use.
batch_type: The training batching stragety to use: can be "examples" or
"tokens".
batch_multiplier: The batch size multiplier to prepare splitting accross
replicated graph parts.
batch_size_multiple: When :obj:`batch_type` is "tokens", ensure that the
resulting batch size is a multiple of this value.
shuffle_buffer_size: The number of elements from which to sample.
length_bucket_width: The width of the length buckets to select batch
candidates from (for efficiency). Set ``None`` to not constrain batch
formation.
maximum_features_length: The maximum length or list of maximum lengths of
the features sequence(s). ``None`` to not constrain the length.
maximum_labels_length: The maximum length of the labels sequence.
``None`` to not constrain the length.
single_pass: If ``True``, makes a single pass over the training data.
num_shards: The number of data shards (usually the number of workers in a
distributed setting).
shard_index: The shard index this data pipeline should read from.
num_threads: The number of elements processed in parallel.
prefetch_buffer_size: The number of batches to prefetch asynchronously. If
``None``, use an automatically tuned value.
cardinality_multiple: Ensure that the dataset cardinality is a multiple of
this value when :obj:`single_pass` is ``True``.
weights: An optional list of weights to create a weighted dataset out of
multiple training files.
Returns:
A ``tf.data.Dataset``.
See Also:
:func:`opennmt.data.training_pipeline`
"""
dataset = self.make_dataset([features_file, labels_file],
training=True)
if weights is not None:
dataset = (dataset, weights)
transform_fns = []
map_func = lambda *arg: self.make_features(element=arg, training=True)
transform_fns.append(lambda dataset: dataset.map(
map_func, num_parallel_calls=num_threads or 4))
transform_fns.append(
dataset_util.filter_examples_by_length(
maximum_features_length=maximum_features_length,
maximum_labels_length=maximum_labels_length,
features_length_fn=self.features_inputter.get_length,
labels_length_fn=self.labels_inputter.get_length))
dataset = dataset_util.training_pipeline(
batch_size,
batch_type=batch_type,
batch_multiplier=batch_multiplier,
batch_size_multiple=batch_size_multiple,
transform_fns=transform_fns,
length_bucket_width=length_bucket_width,
single_pass=single_pass,
num_shards=num_shards,
shard_index=shard_index,
num_threads=num_threads,
shuffle_buffer_size=shuffle_buffer_size,
prefetch_buffer_size=prefetch_buffer_size,
cardinality_multiple=cardinality_multiple)(dataset)
return dataset