python - Google Colab 中的 Keras 调谐器和 TPU

Question

我对 keras 调谐器和 tpu 有一些问题。当我运行下面的代码时，一切正常，网络训练速度很快。

vocab_size = 5000
embedding_dim = 64
max_length = 2000

def create_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Embedding(vocab_size, embedding_dim),
      tf.keras.layers.LSTM(100, dropout=0.5, recurrent_dropout=0.5),
      tf.keras.layers.Dense(embedding_dim, activation='relu'),
      tf.keras.layers.Dense(4, activation='softmax')
  ])
  return model

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

model.fit(train_padded, y_train,
          epochs=10,
          validation_split=0.15,
          verbose=1, batch_size=128)

当我使用 keras tuner 时，神经网络学习缓慢。我相信没有使用TPU。

vocab_size = 5000
max_length = 2000
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

def build_model(hp):
    model = tf.keras.Sequential()
    activation_choice = hp.Choice('activation', values=['relu', 'sigmoid', 'tanh', 'elu', 'selu'])
    embedding_dim = hp.Int('units_hidden', min_value=128, max_value=24, step=8)
    model.add(tf.keras.layers.Embedding(vocab_size, embedding_dim))
    model.add(tf.keras.layers.LSTM(hp.Int('LSTM_Units', min_value=50, max_value=500, step=10), 
                                  dropout=hp.Float('dropout', 0, 0.5, step=0.1, default=0), 
                                  recurrent_dropout=hp.Float('recurrent_dropout', 0, 0.5, step=0.1, default=0)))
    model.add(tf.keras.layers.Dense(embedding_dim, activation=activation_choice))
    model.add(tf.keras.layers.Dense(4, activation='softmax'))
    model.compile(
        optimizer=hp.Choice('optimizer', values=['adam', 'rmsprop', 'SGD']),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=['sparse_categorical_accuracy'])
    return model


with strategy.scope():
    tuner = Hyperband(
      build_model,
      objective='val_accuracy',
      max_epochs=10,
      hyperband_iterations=2)
    tuner.search(train_padded, y_train,
             batch_size=128,
             epochs=10,
             callbacks=[EarlyStopping(patience=1)],
             validation_split=0.15,
             verbose=1)

best_models = tuner.get_best_models(1)
best_model.save('/content/drive/My Drive/best_model.h5')

Notebook link

如何使 keras 调谐器与 TPU 一起工作？

Answer 1

您需要将其传递给调谐器:

tuner = Hyperband(
      build_model,
      objective='val_accuracy',
      max_epochs=10,
      hyperband_iterations=2,
      distribution_strategy=strategy,)

(并删除 strategy.scope() 部分)