Tensorflow — как разделить пакеты между графическими процессорами для прогнозирования на обученных моделях?

Я использую модели, которые я не создавал, а модифицировал (из этого репозитория https://github.com/GeorgeSeif/Semantic-Segmentation-Suite)

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

Я предполагал, что могу сделать что-то вроде этого:

for i, d in enumerate(['\gpu:0', '\gpu:1']):
    with tf.device(d):
        output = sess.run(network, feed_dict={net_input: image_batch[i]})

Но на самом деле это не распределяет задачи по разным графическим процессорам, а также не вызывает ошибки.

Мой вопрос в том, можно ли выделить разные изображения для разных экземпляров сеанса на отдельных графических процессорах без явного изменения предварительного кода сетевого кода. Я хотел бы избежать запуска двух разных скриптов Python с CUDA_VISIBLE_DEVICES = ...

Есть ли простой способ сделать это?

Насколько я понимаю, определения операций должны быть вложены в блок «with tf.device()», однако при выводе операция — это просто загрузка модели и весов, но если я помещу это в блок «with tf. device ()» я получаю сообщение об ошибке, говорящее о том, что граф уже существует и не может быть определен дважды.


person ttocs167    schedule 04.03.2019    source источник


Ответы (1)


tf.device применяется только при построении графа, а не при его выполнении, поэтому перенос вызовов session.run в контекст устройства ничего не делает.

Вместо этого я рекомендую вам использовать tf replicator или стратегию распространения tf (tf .distribute/tf.contrib.distribute в зависимости от версии tf), в частности, MirroredStrategy.

person Alexandre Passos    schedule 14.03.2019