Keras fit_generator работает очень медленно

Question

У меня есть модель Keras, объявленная со следующим кодом:

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.LSTM(units=50, activation="tanh", return_sequences=False, input_shape=(settings["past_size"], len(indicators_used))))
model.add(tf.keras.layers.Dense(3, activation="softmax"))

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
history = model.fit_generator(
    generator=batch_generator(train_x, train_y),
    steps_per_epoch=n_batches_per_epoch,
    epochs=settings["epochs"],
    workers=5,
    use_multiprocessing=True,
    max_queue_size=10000)

Я экспериментировал с настройками workers, use_multiprocessing и max_queue_size, но безрезультатно. Форма ввода (100000, 500, 27).

Функция Batch Generator выглядит следующим образом:

def batch_generator(x, y):
    while True:
        for i in range(n_batches_per_epoch):
            x_train = []
            y_train = []
            for j in range(settings["past_size"] + settings["batch_size"] * i, settings["past_size"] + (settings["batch_size"] * (i + 1))):
                x_train.append(x.iloc[j - past_size:j].to_numpy())
                y_train.append(y.iloc[j].to_numpy())
            return_x, return_y = np.array(x_train), np.array(y_train)
            yield return_x, return_y

Время выполнения:

• Размер партии 256: 767ms/step
• Размер партии512: 1s/step
• Размер партии 1024: 2s/step

Проблема, с которой я сейчас сталкиваюсь, заключается в том, что процесс обучения Keras невероятно медленный. Одна эпоха занимает около 45 минут. Я не могу использовать model.fit(), потому что данные слишком велики для оперативной памяти.

Мое понимание функциональности batch_generator заключалось в том, что функция подготавливает пакеты и загружает их в GPU / TPU, но это не так. Кажется, это так. Этот код работает в Google Colab с TPU Runtime.

Answer 1

В среде Google colab вам необходимо явно преобразовать модель в версию, совместимую с TPU. Это была моя ошибка, когда я в последний раз работал с Google Colab.

TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
    tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

tpu_model имеет тот же интерфейс, что и model.

Guide: https://medium.com/tensorflow/tf-keras-on-tpus-on-colab-674367932aa0

К сожалению, это не похоже на работу с Sequential api.