Ejecute TensorFlow Federated en GPU con Colab

Question

¿Hay alguna forma de utilizar la GPU proporcionada por Colab para ejecutar las sesiones de entrenamiento de TFF más rápido? El entrenamiento de modelos federados requiere más de 1 hora y parece que usar un tiempo de ejecución de GPU no proporciona ningún beneficio en absoluto.

La página TFF de High-Performance Simulation todavía está vacía y no puedo encontrar ninguna guía para usar la GPU con TFF.

¿Cualquier sugerencia? ¡Gracias!

Versiones tf y tff:

2.4.0-dev20200917 
0.16.1

Número de clientes en cada ronda:

70

Especificación del elemento de datos de entrada:

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 manera similar al tutorial de generación de texto, estoy trabajando con secuencia de lugares, el modelo es bastante 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)

Función para crear el modelo:

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

Promedio federado

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

Estado init:

state = fed_avg.initialize()

Bucle de entrenamiento:

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

Tenga en cuenta que este modelo está realizando 0 clientes * 13 pasos de SGD por ronda (cerca de 1,000), aunque una hora todavía parece muy larga. 70 clientes en una sola máquina están superando los límites de la simulación, cuando el número aumenta mucho, comenzamos a mirar configuraciones de múltiples máquinas utilizando el ejecutor remoto.

Algunas cosas para investigar:

• ¿Está vinculada la E / S de simulación ? ¿Qué tan rápido puede iterar el entorno de Python sobre un solo conjunto de datos de cliente? En TF for batch in dataset:y el tiempo, el tiempo que tarda podría ser útil aquí.
• ¿Está limitado el cálculo de la simulación ? Tal vez observe la utilización de la CPU y la GPU. ¿Cuánto tiempo se tarda en ejecutarse keras_model.fit()en un solo conjunto de datos de cliente? La simulación TFF está haciendo esto aproximadamente 70 veces por ronda (una vez para cada cliente).