Tensorflow: Как написать чистый код для параллелизма нескольких графических процессоров?

Question

В настоящее время я реализую большую пользовательскую модель и ссылаюсь на пример с несколькими gpu CIFAR 10, который поставляется вместе с тензорным потоком. Однако код, который я в конечном итоге написал на основе этого, был не чистым и подвержен ошибкам. Напр. Я должен был найти каждую обучаемую переменную и добавить «с tf.device ('/ cpu: 0')».

Существуют ли более эффективные/более чистые способы адаптации для выполнения нескольких gpu?

Большое спасибо за любую поддержку.

Answer 1

Вот example из Рафал

Вы делаете петлю на башнях с телом строящейся i-й башни в with tf.device(assign_to_gpu(i)). Функция assign_to_gpu обрабатывает переменные по-разному и назначает их на «ps-device».

Примечание: мы обнаружили, что, когда чипы p2p подключены, обучение было быстрее, когда переменные были сохранены gpu:0, а не cpu:0