Uruchom TensorFlow Federated na GPU z Colab

Question

Czy istnieje sposób na wykorzystanie GPU dostarczonego przez Colab do szybszego prowadzenia sesji treningowych TFF ? Szkolenie Federated Models trwa dłużej niż 1 godzinę i wydaje się, że korzystanie ze środowiska wykonawczego GPU nie zapewnia żadnych korzyści.

Strona TFF w High-Performance Simulation jest wciąż pusta i nie mogę znaleźć żadnego przewodnika, jak używać GPU z TFF.

Jakieś sugestie? Dziękuję Ci!

wersje tf i tff:

2.4.0-dev20200917 
0.16.1

Liczba klientów w każdej rundzie:

70

Specyfikacja elementu danych wejściowych:

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

Podobnie jak w samouczku dotyczącym generowania tekstu, w którym pracuję z sekwencją miejsc, model jest dość podobny:

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

Funkcja tworzenia modelu:

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

Federacyjna śr

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

Stan init:

state = fed_avg.initialize()

Pętla treningowa:

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

Należy zauważyć, że ten model obsługuje 0 klientów * 13 kroków SGD na rundę (blisko 1000), chociaż jedna godzina nadal wydaje się bardzo długa. 70 klientów na jednej maszynie przesuwa granice symulacji, gdy liczba rośnie znacznie, zaczynamy patrzeć na konfiguracje wielu maszyn za pomocą zdalnego executora.

Kilka rzeczy do zbadania:

• Czy symulacja I / O jest ograniczona? Jak szybko środowisko Python może iterować po jednym zestawie danych klienta? W TF for batch in dataset:i czasie, jak długo to trwa, może być tutaj przydatne.
• Czy obliczenia symulacji są ograniczone? Może obserwuj wykorzystanie procesora i karty graficznej. Ile czasu zajmuje uruchomienie keras_model.fit()na pojedynczym zestawie danych klienta? Symulacja TFF robi to z grubsza 70 razy na rundę (raz dla każdego klienta).