Exécutez TensorFlow Federated sur GPU avec Colab

Question

Existe-t-il un moyen d' utiliser le GPU fourni par Colab pour exécuter les sessions de formation de TFF plus rapidement? La formation des modèles fédérés nécessite plus d'une heure et il semble que l'utilisation d'un moteur d'exécution GPU ne présente aucun avantage.

La page TFF de la simulation haute performance est toujours vide et je ne trouve aucun guide pour utiliser le GPU avec TFF.

Toute suggestion? Merci!

Versions tf et tff:

2.4.0-dev20200917 
0.16.1

Nombre de clients à chaque tour:

70

Spécification de l'élément de données d'entrée:

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))])

De même que le didacticiel de génération de texte dans lequel je travaille avec une séquence de lieux, le modèle est assez similaire:

    # 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)

Fonction pour créer le modèle:

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(),

Moy. Fédérée

# 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()])

État init:

state = fed_avg.initialize()

Boucle d'entraînement:

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']))

Answer 1

A noter que ce modèle effectue 0 client * 13 étapes de SGD par tour (près de 1000), bien qu'une heure semble encore très longue. 70 clients sur une seule machine repoussent les limites de la simulation, lorsque le nombre augmente beaucoup plus, nous commençons à examiner les configurations multi-machines à l'aide de l'exécuteur distant.

Quelques éléments à étudier:

• Les E / S de simulation sont -elles liées? À quelle vitesse l'environnement Python peut-il itérer sur un seul ensemble de données client? Dans TF for batch in dataset:et dans le temps, le temps qu'il faut peut être utile ici.
• Le calcul de la simulation est-il lié? Surveillez peut-être l'utilisation du CPU et du GPU. Combien de temps faut-il pour s'exécuter keras_model.fit()sur un seul ensemble de données client? La simulation TFF fait environ 70 fois par tour (une fois pour chaque client).