def create_character_prediction_task_from_datasets(
train_client_spec: client_spec.ClientSpec,
eval_client_spec: Optional[client_spec.ClientSpec],
sequence_length: int,
train_data: client_data.ClientData,
test_data: client_data.ClientData,
) -> baseline_task.BaselineTask:
"""Creates a baseline task for next-character prediction on Shakespeare.
Args:
train_client_spec: A `tff.simulation.baselines.ClientSpec` specifying how to
preprocess train client data.
eval_client_spec: An optional `tff.simulation.baselines.ClientSpec`
specifying how to preprocess evaluation client data. If set to `None`, the
evaluation datasets will use a batch size of 64 with no extra
preprocessing.
sequence_length: A positive integer dictating the length of each example in
a client's dataset. By default, this is set to
`tff.simulation.baselines.shakespeare.DEFAULT_SEQUENCE_LENGTH`.
train_data: A `tff.simulation.datasets.ClientData` used for training.
test_data: A `tff.simulation.datasets.ClientData` used for testing.
Returns:
A `tff.simulation.baselines.BaselineTask`.
"""
if sequence_length < 1:
raise ValueError('sequence_length must be a positive integer')
if eval_client_spec is None:
eval_client_spec = client_spec.ClientSpec(
num_epochs=1, batch_size=32, shuffle_buffer_size=1)
train_preprocess_fn = char_prediction_preprocessing.create_preprocess_fn(
train_client_spec, sequence_length)
eval_preprocess_fn = char_prediction_preprocessing.create_preprocess_fn(
eval_client_spec, sequence_length)
task_datasets = task_data.BaselineTaskDatasets(
train_data=train_data,
test_data=test_data,
validation_data=None,
train_preprocess_fn=train_preprocess_fn,
eval_preprocess_fn=eval_preprocess_fn)
pad_token, _, _, _ = char_prediction_preprocessing.get_special_tokens()
def model_fn() -> model.Model:
return keras_utils.from_keras_model(
keras_model=char_prediction_models.create_recurrent_model(
vocab_size=VOCAB_LENGTH, sequence_length=sequence_length),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
input_spec=task_datasets.element_type_structure,
metrics=[
keras_metrics.NumTokensCounter(masked_tokens=[pad_token]),
keras_metrics.MaskedCategoricalAccuracy(masked_tokens=[pad_token])
])
return baseline_task.BaselineTask(task_datasets, model_fn)
def create_tag_prediction_task_from_datasets(
train_client_spec: client_spec.ClientSpec,
eval_client_spec: Optional[client_spec.ClientSpec],
word_vocab: List[str],
tag_vocab: List[str],
train_data: client_data.ClientData,
test_data: client_data.ClientData,
validation_data: client_data.ClientData,
) -> baseline_task.BaselineTask:
"""Creates a baseline task for tag prediction on Stack Overflow.
Args:
train_client_spec: A `tff.simulation.baselines.ClientSpec` specifying how to
preprocess train client data.
eval_client_spec: An optional `tff.simulation.baselines.ClientSpec`
specifying how to preprocess evaluation client data. If set to `None`, the
evaluation datasets will use a batch size of 64 with no extra
preprocessing.
word_vocab: A list of strings used for the task's word vocabulary.
tag_vocab: A list of strings used for the task's tag vocabulary.
train_data: A `tff.simulation.datasets.ClientData` used for training.
test_data: A `tff.simulation.datasets.ClientData` used for testing.
validation_data: A `tff.simulation.datasets.ClientData` used for validation.
Returns:
A `tff.simulation.baselines.BaselineTask`.
"""
if eval_client_spec is None:
eval_client_spec = client_spec.ClientSpec(num_epochs=1,
batch_size=100,
shuffle_buffer_size=1)
word_vocab_size = len(word_vocab)
tag_vocab_size = len(tag_vocab)
train_preprocess_fn = tag_prediction_preprocessing.create_preprocess_fn(
train_client_spec, word_vocab, tag_vocab)
eval_preprocess_fn = tag_prediction_preprocessing.create_preprocess_fn(
eval_client_spec, word_vocab, tag_vocab)
task_datasets = task_data.BaselineTaskDatasets(
train_data=train_data,
test_data=test_data,
validation_data=validation_data,
train_preprocess_fn=train_preprocess_fn,
eval_preprocess_fn=eval_preprocess_fn)
def model_fn() -> model.Model:
return keras_utils.from_keras_model(
keras_model=_build_logistic_regression_model(
input_size=word_vocab_size, output_size=tag_vocab_size),
loss=tf.keras.losses.BinaryCrossentropy(
from_logits=False, reduction=tf.keras.losses.Reduction.SUM),
input_spec=task_datasets.element_type_structure,
metrics=[
tf.keras.metrics.Precision(name='precision'),
tf.keras.metrics.Recall(top_k=5, name='recall_at_5'),
])
return baseline_task.BaselineTask(task_datasets, model_fn)
def create_character_recognition_task_from_datasets(
train_client_spec: client_spec.ClientSpec,
eval_client_spec: Optional[client_spec.ClientSpec],
model_id: Union[str, CharacterRecognitionModel], only_digits: bool,
train_data: client_data.ClientData,
test_data: client_data.ClientData) -> baseline_task.BaselineTask:
"""Creates a baseline task for character recognition on EMNIST.
Args:
train_client_spec: A `tff.simulation.baselines.ClientSpec` specifying how to
preprocess train client data.
eval_client_spec: An optional `tff.simulation.baselines.ClientSpec`
specifying how to preprocess evaluation client data. If set to `None`, the
evaluation datasets will use a batch size of 64 with no extra
preprocessing.
model_id: A string identifier for a character recognition model. Must be one
of 'cnn_dropout', 'cnn', or '2nn'. These correspond respectively to a CNN
model with dropout, a CNN model with no dropout, and a densely connected
network with two hidden layers of width 200.
only_digits: A boolean indicating whether to use the full EMNIST-62 dataset
containing 62 alphanumeric classes (`True`) or the smaller EMNIST-10
dataset with only 10 numeric classes (`False`).
train_data: A `tff.simulation.datasets.ClientData` used for training.
test_data: A `tff.simulation.datasets.ClientData` used for testing.
Returns:
A `tff.simulation.baselines.BaselineTask`.
"""
emnist_task = 'character_recognition'
if eval_client_spec is None:
eval_client_spec = client_spec.ClientSpec(
num_epochs=1, batch_size=64, shuffle_buffer_size=1)
train_preprocess_fn = emnist_preprocessing.create_preprocess_fn(
train_client_spec, emnist_task=emnist_task)
eval_preprocess_fn = emnist_preprocessing.create_preprocess_fn(
eval_client_spec, emnist_task=emnist_task)
task_datasets = task_data.BaselineTaskDatasets(
train_data=train_data,
test_data=test_data,
validation_data=None,
train_preprocess_fn=train_preprocess_fn,
eval_preprocess_fn=eval_preprocess_fn)
def model_fn() -> model.Model:
return keras_utils.from_keras_model(
keras_model=_get_character_recognition_model(model_id, only_digits),
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
input_spec=task_datasets.element_type_structure,
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
return baseline_task.BaselineTask(task_datasets, model_fn)
def create_autoencoder_task_from_datasets(
train_client_spec: client_spec.ClientSpec,
eval_client_spec: Optional[client_spec.ClientSpec],
train_data: client_data.ClientData,
test_data: client_data.ClientData) -> baseline_task.BaselineTask:
"""Creates a baseline task for autoencoding on EMNIST.
Args:
train_client_spec: A `tff.simulation.baselines.ClientSpec` specifying how to
preprocess train client data.
eval_client_spec: An optional `tff.simulation.baselines.ClientSpec`
specifying how to preprocess evaluation client data. If set to `None`, the
evaluation datasets will use a batch size of 64 with no extra
preprocessing.
train_data: A `tff.simulation.datasets.ClientData` used for training.
test_data: A `tff.simulation.datasets.ClientData` used for testing.
Returns:
A `tff.simulation.baselines.BaselineTask`.
"""
emnist_task = 'autoencoder'
if eval_client_spec is None:
eval_client_spec = client_spec.ClientSpec(num_epochs=1,
batch_size=64,
shuffle_buffer_size=1)
train_preprocess_fn = emnist_preprocessing.create_preprocess_fn(
train_client_spec, emnist_task=emnist_task)
eval_preprocess_fn = emnist_preprocessing.create_preprocess_fn(
eval_client_spec, emnist_task=emnist_task)
task_datasets = task_data.BaselineTaskDatasets(
train_data=train_data,
test_data=test_data,
validation_data=None,
train_preprocess_fn=train_preprocess_fn,
eval_preprocess_fn=eval_preprocess_fn)
def model_fn() -> model.Model:
return keras_utils.from_keras_model(
keras_model=emnist_models.create_autoencoder_model(),
loss=tf.keras.losses.MeanSquaredError(),
input_spec=task_datasets.element_type_structure,
metrics=[
tf.keras.metrics.MeanSquaredError(),
tf.keras.metrics.MeanAbsoluteError()
])
return baseline_task.BaselineTask(task_datasets, model_fn)
def test_raises_on_different_data_and_model_spec(self):
baseline_task_spec_data = create_task_data()
model_input_spec = (
tf.TensorSpec(shape=(None, 4, 2), dtype=tf.float32, name=None),
tf.TensorSpec(shape=(None, 4, 2), dtype=tf.float32, name=None),
)
inputs = tf.keras.layers.Input(shape=(4, 2))
outputs = tf.keras.layers.Dense(2,
kernel_initializer='ones',
use_bias=False)(inputs)
keras_model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
def tff_model_fn():
return keras_utils.from_keras_model(
keras_model,
loss=tf.keras.losses.MeanSquaredError(),
input_spec=model_input_spec)
with self.assertRaisesRegex(
ValueError,
'Dataset element spec and model input spec do not match'):
baseline_task.BaselineTask(baseline_task_spec_data, tff_model_fn)
def create_image_classification_task_with_datasets(
train_client_spec: client_spec.ClientSpec,
eval_client_spec: Optional[client_spec.ClientSpec],
model_id: Union[str, ResnetModel],
crop_height: int,
crop_width: int,
distort_train_images: bool,
train_data: client_data.ClientData,
test_data: client_data.ClientData,
) -> baseline_task.BaselineTask:
"""Creates a baseline task for image classification on CIFAR-100.
Args:
train_client_spec: A `tff.simulation.baselines.ClientSpec` specifying how to
preprocess train client data.
eval_client_spec: An optional `tff.simulation.baselines.ClientSpec`
specifying how to preprocess evaluation client data. If set to `None`, the
evaluation datasets will use a batch size of 64 with no extra
preprocessing.
model_id: A string identifier for a digit recognition model. Must be one of
`resnet18`, `resnet34`, `resnet50`, `resnet101` and `resnet152. These
correspond to various ResNet architectures. Unlike standard ResNet
architectures though, the batch normalization layers are replaced with
group normalization.
crop_height: An integer specifying the desired height for cropping images.
Must be between 1 and 32 (the height of uncropped CIFAR-100 images). By
default, this is set to
`tff.simulation.baselines.cifar100.DEFAULT_CROP_HEIGHT`.
crop_width: An integer specifying the desired width for cropping images.
Must be between 1 and 32 (the width of uncropped CIFAR-100 images). By
default this is set to
`tff.simulation.baselines.cifar100.DEFAULT_CROP_WIDTH`.
distort_train_images: Whether to distort images in the train preprocessing
function.
train_data: A `tff.simulation.datasets.ClientData` used for training.
test_data: A `tff.simulation.datasets.ClientData` used for testing.
Returns:
A `tff.simulation.baselines.BaselineTask`.
"""
if crop_height < 1 or crop_width < 1 or crop_height > 32 or crop_width > 32:
raise ValueError(
'The crop_height and crop_width must be between 1 and 32.')
crop_shape = (crop_height, crop_width, 3)
if eval_client_spec is None:
eval_client_spec = client_spec.ClientSpec(num_epochs=1,
batch_size=64,
shuffle_buffer_size=1)
train_preprocess_fn = image_classification_preprocessing.create_preprocess_fn(
train_client_spec,
crop_shape=crop_shape,
distort_image=distort_train_images)
eval_preprocess_fn = image_classification_preprocessing.create_preprocess_fn(
eval_client_spec, crop_shape=crop_shape)
task_datasets = task_data.BaselineTaskDatasets(
train_data=train_data,
test_data=test_data,
validation_data=None,
train_preprocess_fn=train_preprocess_fn,
eval_preprocess_fn=eval_preprocess_fn)
def model_fn() -> model.Model:
return keras_utils.from_keras_model(
keras_model=_get_resnet_model(model_id, crop_shape),
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
input_spec=task_datasets.element_type_structure,
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
return baseline_task.BaselineTask(task_datasets, model_fn)
def test_construct_raises_on_non_tff_model(self):
with self.assertRaisesRegex(
TypeError, 'Expected model_fn to output a tff.learning.Model'):
baseline_task.BaselineTask(create_task_data(), keras_model_builder)
def test_construct_raises_on_non_callable_model_fn(self):
with self.assertRaises(attr.exceptions.NotCallableError):
baseline_task.BaselineTask(create_task_data(), 'bad_input')
def test_construct_raises_on_non_baseline_task_spec_datasets(self):
with self.assertRaises(TypeError):
baseline_task.BaselineTask(create_dataset(0), tff_model_builder)
def test_constructs_with_matching_dataset_and_model_spec(self):
baseline_task_spec = baseline_task.BaselineTask(
create_task_data(), tff_model_builder)
self.assertIsInstance(baseline_task_spec, baseline_task.BaselineTask)
def create_word_prediction_task_from_datasets(
train_client_spec: client_spec.ClientSpec,
eval_client_spec: Optional[client_spec.ClientSpec],
sequence_length: int,
vocab_size: int,
num_out_of_vocab_buckets: int,
train_data: client_data.ClientData,
test_data: client_data.ClientData,
validation_data: client_data.ClientData,
) -> baseline_task.BaselineTask:
"""Creates a baseline task for next-word prediction on Stack Overflow.
The goal of the task is to take `sequence_length` words from a post and
predict the next word. Here, all posts are drawn from the Stack Overflow
forum, and a client corresponds to a user.
Args:
train_client_spec: A `tff.simulation.baselines.ClientSpec` specifying how to
preprocess train client data.
eval_client_spec: An optional `tff.simulation.baselines.ClientSpec`
specifying how to preprocess evaluation client data. If set to `None`, the
evaluation datasets will use a batch size of 64 with no extra
preprocessing.
sequence_length: A positive integer dictating the length of each word
sequence in a client's dataset. By default, this is set to
`tff.simulation.baselines.stackoverflow.DEFAULT_SEQUENCE_LENGTH`.
vocab_size: Integer dictating the number of most frequent words in the
entire corpus to use for the task's vocabulary. By default, this is set to
`tff.simulation.baselines.stackoverflow.DEFAULT_WORD_VOCAB_SIZE`.
num_out_of_vocab_buckets: The number of out-of-vocabulary buckets to use.
train_data: A `tff.simulation.datasets.ClientData` used for training.
test_data: A `tff.simulation.datasets.ClientData` used for testing.
validation_data: A `tff.simulation.datasets.ClientData` used for validation.
Returns:
A `tff.simulation.baselines.BaselineTask`.
"""
if sequence_length < 1:
raise ValueError('sequence_length must be a positive integer')
if vocab_size < 1:
raise ValueError('vocab_size must be a positive integer')
if num_out_of_vocab_buckets < 1:
raise ValueError('num_out_of_vocab_buckets must be a positive integer')
vocab = list(stackoverflow.load_word_counts(vocab_size=vocab_size).keys())
if eval_client_spec is None:
eval_client_spec = client_spec.ClientSpec(num_epochs=1,
batch_size=64,
shuffle_buffer_size=1)
train_preprocess_fn = word_prediction_preprocessing.create_preprocess_fn(
train_client_spec,
vocab,
sequence_length=sequence_length,
num_out_of_vocab_buckets=num_out_of_vocab_buckets)
eval_preprocess_fn = word_prediction_preprocessing.create_preprocess_fn(
eval_client_spec,
vocab,
sequence_length=sequence_length,
num_out_of_vocab_buckets=num_out_of_vocab_buckets)
task_datasets = task_data.BaselineTaskDatasets(
train_data=train_data,
test_data=test_data,
validation_data=validation_data,
train_preprocess_fn=train_preprocess_fn,
eval_preprocess_fn=eval_preprocess_fn)
special_tokens = word_prediction_preprocessing.get_special_tokens(
vocab_size, num_out_of_vocab_buckets=num_out_of_vocab_buckets)
pad_token = special_tokens.padding
oov_tokens = special_tokens.out_of_vocab
eos_token = special_tokens.end_of_sentence
def metrics_builder():
return [
keras_metrics.NumTokensCounter(masked_tokens=[pad_token]),
keras_metrics.MaskedCategoricalAccuracy(name='accuracy',
masked_tokens=[pad_token]),
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_without_out_of_vocab',
masked_tokens=[pad_token] + oov_tokens),
# Notice that the beginning of sentence token never appears in the
# ground truth label.
keras_metrics.MaskedCategoricalAccuracy(
name='accuracy_without_out_of_vocab_or_end_of_sentence',
masked_tokens=[pad_token, eos_token] + oov_tokens),
]
# The total vocabulary size is the number of words in the vocabulary, plus
# the number of out-of-vocabulary tokens, plus three tokens used for
# padding, beginning of sentence and end of sentence.
extended_vocab_size = (vocab_size +
special_tokens.get_number_of_special_tokens())
def model_fn() -> model.Model:
return keras_utils.from_keras_model(
keras_model=word_prediction_models.create_recurrent_model(
vocab_size=extended_vocab_size),
loss=tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True),
input_spec=task_datasets.element_type_structure,
metrics=metrics_builder())
return baseline_task.BaselineTask(task_datasets, model_fn)
def create_tag_prediction_task_from_datasets(
train_client_spec: client_spec.ClientSpec,
eval_client_spec: Optional[client_spec.ClientSpec],
word_vocab_size: int,
tag_vocab_size: int,
train_data: client_data.ClientData,
test_data: client_data.ClientData,
validation_data: client_data.ClientData,
) -> baseline_task.BaselineTask:
"""Creates a baseline task for tag prediction on Stack Overflow.
Args:
train_client_spec: A `tff.simulation.baselines.ClientSpec` specifying how to
preprocess train client data.
eval_client_spec: An optional `tff.simulation.baselines.ClientSpec`
specifying how to preprocess evaluation client data. If set to `None`, the
evaluation datasets will use a batch size of 64 with no extra
preprocessing.
word_vocab_size: Integer dictating the number of most frequent words in the
entire corpus to use for the task's vocabulary. By default, this is set to
`tff.simulation.baselines.stackoverflow.DEFAULT_WORD_VOCAB_SIZE`.
tag_vocab_size: Integer dictating the number of most frequent tags in the
entire corpus to use for the task's labels. By default, this is set to
`tff.simulation.baselines.stackoverflow.DEFAULT_TAG_VOCAB_SIZE`.
train_data: A `tff.simulation.datasets.ClientData` used for training.
test_data: A `tff.simulation.datasets.ClientData` used for testing.
validation_data: A `tff.simulation.datasets.ClientData` used for validation.
Returns:
A `tff.simulation.baselines.BaselineTask`.
"""
if word_vocab_size < 1:
raise ValueError('word_vocab_size must be a positive integer')
if tag_vocab_size < 1:
raise ValueError('tag_vocab_size must be a positive integer')
word_vocab = list(stackoverflow.load_word_counts(vocab_size=word_vocab_size))
tag_vocab = list(stackoverflow.load_tag_counts().keys())[:tag_vocab_size]
if eval_client_spec is None:
eval_client_spec = client_spec.ClientSpec(
num_epochs=1, batch_size=100, shuffle_buffer_size=1)
train_preprocess_fn = tag_prediction_preprocessing.create_preprocess_fn(
train_client_spec, word_vocab, tag_vocab)
eval_preprocess_fn = tag_prediction_preprocessing.create_preprocess_fn(
eval_client_spec, word_vocab, tag_vocab)
task_datasets = task_data.BaselineTaskDatasets(
train_data=train_data,
test_data=test_data,
validation_data=validation_data,
train_preprocess_fn=train_preprocess_fn,
eval_preprocess_fn=eval_preprocess_fn)
def model_fn() -> model.Model:
return keras_utils.from_keras_model(
keras_model=_build_logistic_regression_model(
input_size=word_vocab_size, output_size=tag_vocab_size),
loss=tf.keras.losses.BinaryCrossentropy(
from_logits=False, reduction=tf.keras.losses.Reduction.SUM),
input_spec=task_datasets.element_type_structure,
metrics=[
tf.keras.metrics.Precision(name='precision'),
tf.keras.metrics.Recall(top_k=5, name='recall_at_5'),
])
return baseline_task.BaselineTask(task_datasets, model_fn)