Запустите TensorFlow Federated на GPU с Colab

Question

Есть ли способ использовать графический процессор, предоставленный Colab, для ускорения тренировок TFF ? Для обучения федеративных моделей требуется более 1 часа, и кажется, что использование среды выполнения GPU не дает никаких преимуществ.

Страница TFF высокопроизводительного моделирования все еще пуста, и я не могу найти никакого руководства по использованию графического процессора с TFF.

Любое предложение? Спасибо!

Версии tf и tff:

2.4.0-dev20200917 
0.16.1

Количество клиентов в каждом раунде:

70

Спецификация элемента входных данных:

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

Как и в учебнике по созданию текста, я работаю с последовательностью мест, модель очень похожа:

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

Функция для создания модели:

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

Федеративная средняя

# 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 = fed_avg.initialize()

Цикл обучения:

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

Отметим, что эта модель выполняет 0 клиентов * 13 шагов SGD за раунд (около 1000), хотя один час все еще кажется очень длинным. 70 клиентов на одной машине раздвигают границы моделирования, и когда их число становится намного выше, мы начинаем смотреть на настройки нескольких машин с помощью удаленного исполнителя.

Некоторые вещи для расследования:

• Ограничен ли ввод / вывод моделирования ? Насколько быстро среда Python может перебирать один набор данных клиента? В TF for batch in dataset:и времени может быть полезно здесь.
• Ограничены ли вычислительные возможности моделирования ? Может быть, посмотреть загрузку CPU и GPU. Сколько времени нужно для работы keras_model.fit()с одним набором данных клиента? Моделирование TFF делает это примерно 70 раз за раунд (один раз для каждого клиента).