Распределение модели Keras по нескольким графическим процессорам

Question

Я пытаюсь создать очень большую модель Keras и распределить ее по нескольким графическим процессорам.Чтобы было ясно, я не пытаюсь поместить несколько копий одной и той же модели на несколько графических процессоров;Я пытаюсь поставить одну большую модель на несколько графических процессоров.Я использовал функцию multi_gpu_model в Keras, но, основываясь на множестве ошибок нехватки памяти, которые я получил, делая это, похоже, что она просто копирует модель, а не распределяет ее, как хотелось бы.

Я заглянул в Horovod, но из-за того, что у меня запущено много инструментов ведения журналов для окон, я не решаюсь его использовать.

Похоже, что мне остается использовать только tf.estimators.Хотя из документации не ясно, как я мог бы использовать эти оценки, чтобы делать то, что я пытаюсь сделать.Например, какая стратегия распространения в tf.contrib.distribute позволила бы мне эффективно распределить модель по тому пути, который я собираюсь сделать?

Возможно ли то, что я пытаюсь сделать с оценщиками, и если да, токакую стратегию я должен использовать?

Answer 1

Вы можете использовать Estimator API.Конвертируйте вашу модель, используя tf.keras.estimator.model_to_estimator

session_config = tf.ConfigProto(allow_soft_placement=True)
distribute = tf.contrib.distribute.MirroredStrategy(num_gpus=4)
run_config = tf.estimator.RunConfig(train_distribute=distribute)
your_network = tf.keras.estimator.model_to_estimator(model_fn=your_keras_model, config=run_config)
your_network.train(input_fn)

Не забудьте скомпилировать модель

Answer 2

Вам нужен параллелизм устройства. В этом разделе FAQ по Keras приведен пример, как это сделать с Keras:

# Model where a shared LSTM is used to encode two different sequences in parallel
input_a = keras.Input(shape=(140, 256))
input_b = keras.Input(shape=(140, 256))

shared_lstm = keras.layers.LSTM(64)

# Process the first sequence on one GPU
with tf.device_scope('/gpu:0'):
    encoded_a = shared_lstm(tweet_a)
# Process the next sequence on another GPU
with tf.device_scope('/gpu:1'):
    encoded_b = shared_lstm(tweet_b)

# Concatenate results on CPU
with tf.device_scope('/cpu:0'):
    merged_vector = keras.layers.concatenate([encoded_a, encoded_b],
                                             axis=-1)

Answer 3

Вы можете вручную назначить разные части вашей модели Keras различным графическим процессорам, используя серверную часть TensorFlow. В этом руководстве приведены подробные примеры, а в этой статье объясняется использование Keras с TensorFlow.

import tensorflow as tf

with tf.device("/device:GPU:0"):
    #Create first part of your neural network

with tf.device("/device:GPU:1"):
    #Create second part of your neural network

#...

with tf.device("/device:GPU:n"):
    #Create nth part of your neural network

Осторожно: задержки связи между ЦП и несколькими графическими процессорами могут привести к значительным накладным расходам.на обучение.