Is there a way to use the GPU provided by Colab to run the training sessions of TFF faster? Training Federated Models requires more than 1 hour and it seems that using a GPU runtime does not provide any benefit at all.
The TFF page of High-Performance Simulation is still empty and I cannot find any guide to use the GPU with TFF.
Any suggestion? Thank you!
tf and tff versions:
2.4.0-dev20200917
0.16.1
Number of clients at each round:
70
Input data element spec:
OrderedDict([('x',
OrderedDict([('start_place',
TensorSpec(shape=(8, 8), dtype=tf.int32, name=None)),
('start_hour_sin',
TensorSpec(shape=(8, 8), dtype=tf.float64, name=None)),
('start_hour_cos',
TensorSpec(shape=(8, 8), dtype=tf.float64, name=None)),
('week_day_sin',
TensorSpec(shape=(8, 8), dtype=tf.float64, name=None)),
('week_day_cos',
TensorSpec(shape=(8, 8), dtype=tf.float64, name=None)),
('weekend',
TensorSpec(shape=(8, 8), dtype=tf.int32, name=None)),
('month',
TensorSpec(shape=(8, 8), dtype=tf.int32, name=None))])),
('y', TensorSpec(shape=(8, 8), dtype=tf.int32, name=None))])
Similarly to the text generation tutorial i'm working with sequence of places, the model is quite similar:
# Create a model
def create_keras_model(number_of_places, batch_size):
# Shortcut to the layers package
l = tf.keras.layers
# Now we need to define an input dictionary.
# Where the keys are the column names
# This is a model with multiple inputs, so we need to declare and input layer for each feature
feature_inputs = {
'start_hour_sin': tf.keras.Input((N-1, ), batch_size=batch_size, name='start_hour_sin'),
'start_hour_cos': tf.keras.Input((N-1, ), batch_size=batch_size, name='start_hour_cos'),
'weekend': tf.keras.Input((N-1, ), batch_size=batch_size, name='weekend'),
'week_day_sin': tf.keras.Input((N-1, ), batch_size=batch_size, name='week_day_sin'),
'week_day_cos': tf.keras.Input((N-1, ), batch_size=batch_size, name='week_day_cos'),
}
# We cannot use anarray of features as always because we have sequences and we cannot match the shape otherwise
# We have to do one by one
start_hour_sin = feature_column.numeric_column("start_hour_sin", shape=(N-1))
hour_sin_feature = l.DenseFeatures(start_hour_sin)(feature_inputs)
start_hour_cos = feature_column.numeric_column("start_hour_cos", shape=(N-1))
hour_cos_feature = l.DenseFeatures(start_hour_cos)(feature_inputs)
weekend = feature_column.numeric_column("weekend", shape=(N-1))
weekend_feature = l.DenseFeatures(weekend)(feature_inputs)
week_day_sin = feature_column.numeric_column("week_day_sin", shape=(N-1))
week_day_sin_feature = l.DenseFeatures(week_day_sin)(feature_inputs)
week_day_cos = feature_column.numeric_column("week_day_cos", shape=(N-1))
week_day_cos_feature = l.DenseFeatures(week_day_cos)(feature_inputs)
# We have also to add a dimension to then concatenate
hour_sin_feature = tf.expand_dims(hour_sin_feature, -1)
hour_cos_feature = tf.expand_dims(hour_cos_feature, -1)
weekend_feature = tf.expand_dims(weekend_feature, -1)
week_day_sin_feature = tf.expand_dims(week_day_sin_feature, -1)
week_day_cos_feature = tf.expand_dims(week_day_cos_feature, -1)
# Declare the dictionary for the places sequence as before
sequence_input = {
'start_place': tf.keras.Input((N-1,), batch_size=batch_size, dtype=tf.dtypes.int32, name='start_place') # add batch_size=batch_size in case of stateful GRU
}
# Handling the categorical feature sequence using one-hot
places_one_hot = feature_column.sequence_categorical_column_with_vocabulary_list(
'start_place', [i for i in range(number_of_places)])
# Embed the one-hot encoding
places_embed = feature_column.embedding_column(places_one_hot, embedding_dim)
# With an input sequence we can't use the DenseFeature layer, we need to use the SequenceFeatures
sequence_features, sequence_length = tf.keras.experimental.SequenceFeatures(places_embed)(sequence_input)
input_sequence = l.Concatenate(axis=2)([ sequence_features, hour_sin_feature, hour_cos_feature, weekend_feature, week_day_sin_feature, week_day_cos_feature])
# Rnn
recurrent = l.GRU(rnn_units,
batch_size=batch_size, #in case of stateful
return_sequences=True,
dropout=0.5,
stateful=True,
recurrent_initializer='glorot_uniform')(input_sequence)
# Last layer with an output for each places
dense_1 = layers.Dense(number_of_places)(recurrent)
# Softmax output layer
output = l.Softmax()(dense_1)
# To return the Model, we need to define it's inputs and outputs
# In out case, we need to list all the input layers we have defined
inputs = list(feature_inputs.values()) + list(sequence_input.values())
# Return the Model
return tf.keras.Model(inputs=inputs, outputs=output)
Function to create the model:
def create_tff_model():
# TFF uses an `input_spec` so it knows the types and shapes
# that your model expects.
input_spec = preprocessed_example_dataset.element_spec
keras_model_clone = create_keras_model(number_of_places, batch_size=BATCH_SIZE)
return tff.learning.from_keras_model(
keras_model_clone,
input_spec=input_spec,
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
Federated Avg
# This command builds all the TensorFlow graphs and serializes them:
fed_avg = tff.learning.build_federated_averaging_process(
model_fn=create_tff_model,
client_optimizer_fn=lambda: tf.keras.optimizers.Adam(learning_rate=0.001),
server_optimizer_fn=lambda: tf.keras.optimizers.Adam(learning_rate=0.06))
metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
State init:
state = fed_avg.initialize()
Training Loop:
NUM_ROUNDS = 10
for round_num in range(1, NUM_ROUNDS + 1):
print('Round {r}'.format(r=round_num))
state, metrics = fed_avg.next(state, train_data)
train_metrics = metrics['train']
print('\tTrain: loss={l:.3f}, accuracy={a:.3f}'.format(l=train_metrics['loss'], a=train_metrics['sparse_categorical_accuracy']))