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

В качестве предисловия к этому я хотел бы отметить, что ваша модель может не работать быстрее на нескольких графических процессорах, если вы на самом деле не привязаны к графическому процессору; Некоторые случаи, когда это может произойти, включают, когда вы используете генератор с вашими данными, и его создание связано с ЦП / вводом-выводом, или если ваша модель не особенно сложна и вы привязаны к памяти при передаче данных на свой графический процессор.

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

import numpy as np
import keras
from keras.models import Sequential
from keras.layers import Dense
model = Sequential()
model.add(Dense(4000, input_dim=8000, activation='tanh'))
model.add(Dense(2000, input_dim=8000, activation='relu'))
model.add(Dense(500, activation='relu'))
model.add(Dense(300, activation='relu'))
model.add(Dense(1, activation='sigmoid'))
print (model.summary())
model = make_parallel(model, 4)
optimizer = keras.optimizers.Adam(lr=0.0001)
model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
x = np.random.rand(131072, 8000)
y = np.random.randint(0, 2, (131072, 1))
model.fit(x, y, batch_size=2048*4)

Еще одна вещь, которую вам нужно отметить, - это то, что если вы хотите повысить производительность, вы должны умножить обычный размер ваших пакетов на количество используемых вами графических процессоров.

При оценке производительности я обнаружил, что мне нужны довольно большие модели, чтобы заметить разницу между производительностью одного графического процессора и производительностью нескольких графических процессоров; Керас говорит мне, что модель, которую я определил выше, имеет 41 157 101 параметр, а модели, которые были значительно меньше, похоже, не получили большого прироста производительности от настройки нескольких графических процессоров.

Внутренности

make_parallel - относительно простая функция:

  • Он создает копию вашей модели на N графических процессорах, которым вы указываете
  • Он разбивает вашу партию на N меньших партий равного размера.
  • Он передает каждую меньшую партию в соответствующую модель.
  • Он объединяет выходные данные моделей

Часть, которой этого не хватает в большинстве руководств по TensorFlow Multi-GPU, заключается в том, что мы не вычисляем отдельную функцию потерь для каждого графического процессора, а затем усредняем результаты перед применением изменений. У нас есть только одна функция потерь, определенная для нашей модели. Чтобы избежать попадания всех наших градиентов на одно и то же устройство, мы передаем TensorFlow флаг colocate_gradients_with_ops, чтобы попросить его вычислить градиенты на графическом процессоре, на котором выполняется каждая операция, хотя выигрыш в производительности не кажется большим, поэтому неясно, насколько хорошо работает этот подход.