Можно ли установить привязку к графическому процессору для NN со смешанной точностью, когда FP32 и FP16 будут использовать разные графические процессоры?

Question

У меня есть GTX 1080 и RTX 2080. Я хочу тренироваться, используя оба, но так как RTX может обрабатывать FP16 вдвое быстрее, я бы хотел настроить его так, чтобы обучение проходило на нескольких GPU и RTXобрабатывает слои FP16, а GTX - слои FP32.

Возможно ли это при тензорном потоке, pytorch или keras?

Answer 1

Tensorflow

В TF можно указать для каждого слоя, на каком устройстве выполняться (GPU, CPU или определенный GPU, если у вас несколько GPU ...).Это делается с помощью оператора with tf.device('device_name') (необходимо указать значение device_name).См. Использование нескольких графических процессоров раздел.

Keras

Поскольку это возможно в TF, это означает, что вы можете использовать его также в Keras, если вы используете TF в качестве внутреннего интерфейса Keras.(Keras - это просто высокоуровневый API нейронных сетей.)

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