Вопрос или проблема
Я застрял с такой проблемой уже некоторое время. У меня есть настройка AWS с 500 ГБ ОЗУ и около 7 ГПУ. Проблема в том, что каждый раз, когда я пытаюсь запустить свой Keras с TensorFlow как бекенды, у меня заканчивается память. Я также выяснил причину этого. Дело в том, что у каждого ГПУ всего 12 ГБ памяти, в то время как моей модели требуется больше. Как я могу запустить модель так, чтобы она использовала память всех ГПУ вместе, чтобы загрузить модель, а не только полагаться на память одного ГПУ для загрузки всей модели и выхода за пределы памяти? Я пробовал параллелизм модели с Keras, и он, похоже, установлен правильно, так как при печати слоев каждый слой назначен программному ГПУ, но модель все равно пытается загрузиться в память одного ГПУ, т.е. только 11 ГБ и вскоре заканчивается память.
Есть идеи, что происходит?
with tf.device('/gpu:0'):
x = conv2d_bn(img_input, 32, 3, 3, strides=(2, 2), padding='valid')
x = conv2d_bn(x, 32, 3, 3, padding='valid')
x = conv2d_bn(x, 64, 3, 3)
x = MaxPooling2D((3, 3), strides=(2, 2))(x)
x = conv2d_bn(x, 80, 1, 1, padding='valid')
x = conv2d_bn(x, 192, 3, 3, padding='valid')
x = MaxPooling2D((3, 3), strides=(2, 2))(x)
# mixed 0, 1, 2: 35 x 35 x 256
branch1x1 = conv2d_bn(x, 64, 1, 1)
branch5x5 = conv2d_bn(x, 48, 1, 1)
branch5x5 = conv2d_bn(branch5x5, 64, 5, 5)
branch3x3dbl = conv2d_bn(x, 64, 1, 1)
branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
branch_pool = AveragePooling2D((3, 3), strides=(1, 1), padding='same')(x)
branch_pool = conv2d_bn(branch_pool, 32, 1, 1)
x = layers.concatenate(
[branch1x1, branch5x5, branch3x3dbl, branch_pool],
axis=channel_axis,
name="mixed0")
print(x)
with tf.device('/gpu:1'):
# mixed 1: 35 x 35 x 256
branch1x1 = conv2d_bn(x, 64, 1, 1)
branch5x5 = conv2d_bn(x, 48, 1, 1)
branch5x5 = conv2d_bn(branch5x5, 64, 5, 5)
branch3x3dbl = conv2d_bn(x, 64, 1, 1)
branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
branch_pool = AveragePooling2D((3, 3), strides=(1, 1), padding='same')(x)
branch_pool = conv2d_bn(branch_pool, 64, 1, 1)
x = layers.concatenate(
[branch1x1, branch5x5, branch3x3dbl, branch_pool],
axis=channel_axis,
name="mixed1")
# mixed 2: 35 x 35 x 256
branch1x1 = conv2d_bn(x, 64, 1, 1)
branch5x5 = conv2d_bn(x, 48, 1, 1)
branch5x5 = conv2d_bn(branch5x5, 64, 5, 5)
branch3x3dbl = conv2d_bn(x, 64, 1, 1)
branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
branch3x3dbl = conv2d_bn(branch3x3dbl, 96, 3, 3)
branch_pool = AveragePooling2D((3, 3), strides=(1, 1), padding='same')(x)
branch_pool = conv2d_bn(branch_pool, 64, 1, 1)
x = layers.concatenate(
[branch1x1, branch5x5, branch3x3dbl, branch_pool],
axis=channel_axis,
name="mixed2")
print(x)
Редактировать: Вот ссылка на код.
То, что вы описываете, – это параллелизм модели – разные части одной модели работают на разных устройствах, обрабатывая одну партию данных вместе.
Существует множество способов реализации параллелизма модели в Keras. Вы, вероятно, хотите сделать синхронное обучение на одном хосте с несколькими устройствами – одна машина с несколькими ГПУ. Каждое устройство будет запускать копию вашей модели (называемую репликой).
Код выглядит примерно так:
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()
# Открыть область стратегии.
with strategy.scope():
# Все, что создает переменные, должно находиться в области стратегии.
# Как правило, это только построение модели и `compile()`.
model = Model(...)
model.compile(...)
# Обучите модель на всех доступных устройствах.
model.fit(train_dataset, validation_data=val_dataset, ...)
# Проверьте модель на всех доступных устройствах.
model.evaluate(test_dataset)
Ответ или решение
Проблема, с которой вы столкнулись, заключается в ограничениях памяти при работе с большими моделями глубокого обучения на устройствах с ограниченным объемом видеопамяти. Это распространенная ситуация, особенно при использовании моделей с высокой сложностью, таких как Inception v3. В данной статье мы рассмотрим, как решить проблему нехватки памяти за счет использования параллельного моделирования (model parallelism) в Keras и TensorFlow, а также обсудим важные аспекты использования многопроцессорной настройки.
Проблема с параллельным моделированием в Keras и TensorFlow
Вы правильно определили, что каждая GPU в вашем кластере имеет ограничение в 12 ГБ видеопамяти, а ваша модель превышает это ограничение. При использовании воспроизведения устройства TensorFlow с помощью tf.device, как показано в вашем коде, вы указываете, какие слои должны выполняться на каких устройствах. Однако, это не всегда имеет желаемый эффект, особенно если не используется правильная стратегия распределения.
Решение с использованием Стратегии Распределенного Обучения
Вместо ручного распределения слоев модели по устройствам, рекомендовано использовать tf.distribute.Strategy, который значительно упрощает процесс распределенного обучения и помогает избежать проблем с памятью. Этот подход автоматически масштабирует вашу модель и распределяет веса и градиенты между несколькими устройствами.
Пример кода
Вот пример реализации с использованием технологии MirroredStrategy для достижения оптимального параллелизма:
import tensorflow as tf
from tensorflow.keras import layers, Model
# Создаем стратегию
strategy = tf.distribute.MirroredStrategy()
# Открываем область стратегии
with strategy.scope():
# Конструируем модель и компилируем ее
model = Model(inputs=img_input, outputs=x)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
# Обучаем модель на всех доступных устройствах
model.fit(train_dataset, validation_data=val_dataset, epochs=10)
# Оценка модели на всех доступных устройствах
model.evaluate(test_dataset)Ключевые моменты
-
Применение Стратегии: Используя
MirroredStrategy, вы гарантируете, что модель будет распределена по всем доступным GPU, а данные будут обрабатываться параллельно, что увеличивает эффективность использования памяти. -
Область стратегии: Все операции, создающие переменные модели, должны выполняться в области стратегии. Это гарантирует, что переменные будут правильно синхронизированы.
-
Оптимизация использования памяти: Стратегия автоматически делит данные между устройствами, избегая необходимости вручную управлять устройством для каждого слоя, что может привести к ошибкам и плохой производительности.
Заключение
Использование tf.distribute.Strategy, особенно MirroredStrategy, значительно упрощает процесс параллельного обучения и помогает справиться с проблемами, связанными с ограниченной памятью на GPU. Это решение увеличивает эффективность обучения модели и позволяет выполнять сложные задачи глубокого обучения, требующие значительных вычислительных ресурсов.
Если у вас возникнут дополнительные вопросы или потребуется помощь в реализации этой стратегии, пожалуйста, не стесняйтесь обращаться!