Как разбить график (модель) TensorFlow на несколько графических процессоров, чтобы избежать OOM?

Итак, у меня есть очень большая и глубокая модель, которую я реализовал с помощью TensorFlow r1.2, работающей на NVIDIA Tesla k40 с 12 ГБ памяти. Модель состоит из нескольких RNN, набора весовых матриц и матриц вложения, а также векторов смещения. Когда я запустил программу обучения, сначала потребовалось около 2-3 часов, чтобы построить модель, а затем произошел сбой из-за проблем с OOM. Я попытался уменьшить размер пакета до одного образца данных на пакет, но все равно столкнулся с той же проблемой.

Если я tensorflow muitlple gpu в Google, то найденные мною примеры в основном касались использования нескольких графических процессоров с помощью параллельного дизайна модели, что означает, что каждый графический процессор должен запускать одно и то же. график и пусть ЦП вычислит общий градиент, таким образом, распространяется обратно к каждому параметру.

Я знаю, что одним из возможных решений может быть запуск модели на графическом процессоре с большей памятью. Но мне интересно, есть ли способ разделить мой график (модель) на разные части последовательно и назначить их разным графическим процессорам?


tensorflow tensorflow-gpu
person Vincent Huang    schedule 27.07.2017    source источник

Ответы (1)


arrow_upward
1
arrow_downward

Официальное руководство по использованию графических процессоров показывает вам этот пример в разделе «Использование нескольких графических процессоров». Вам просто нужно создать операции в разных контекстах tf.device; узлы по-прежнему будут добавлены в тот же граф, но они будут аннотированы директивами устройства, указывающими, где они должны быть запущены. Например:

with tf.device("/gpu:0"):
    net0 = make_subnet0()
with tf.device("/gpu:1"):
    net1 = make_subnet1()
result = combine_subnets(net0, net1)
person jdehesa    schedule 27.07.2017