Tensorflow 2: размещение переменных и операций на разных устройствах (GPU / CPU) с использованием tf.device и MirroredStrategy

Question

Я хотел бы хранить разные переменные и запускать разные операции тензорного потока на разных устройствах (GPU / CPU). Я хочу сделать это, потому что моя модель не может поместиться на одном графическом процессоре. Пробую тренировать вложения. Я хочу разместить повторяющиеся слова на всех графических процессорах и необщие слова на разных устройствах (в рамках предварительной обработки я все это делаю) и передавать предложения на соответствующие графические процессоры.

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

Ниже приведен фрагмент. Для простоты скажем, что я хочу сделать y = x * w1 + x * w2 + 2 * w3, где w1, w2 и w3 - обучаемые переменные. Ставлю w1 на GPU0 и w2 на GPU1. Вместо того, чтобы выполнять 2 * w3 на одном устройстве, скажите, что я хочу разделить его и выполнить на 2 устройствах (чтобы увидеть, как работает зеркальная стратегия). Я выполняю соответствующие вычисления на соответствующих графических процессорах. Прежде всего, это правильный способ сделать это (использовать зеркальную область видимости для назначения устройств переменным). Во-вторых, если это неверно, почему тензорный поток не выдает никаких ошибок, почему он застревает и почему показывает 100% использование графического процессора?

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

import os

#os.environ['TF2_BEHAVIOR'] = '1'
import tensorflow as tf
tf.compat.v1.enable_eager_execution()

#tf.compat.v1.disable_eager_execution()

tf.config.set_soft_device_placement(False)
tf.debugging.set_log_device_placement(True)
import numpy as np

class ProdLayer(tf.keras.layers.Layer):
    def __init__(self, name):
        super(ProdLayer, self).__init__()
        self.w = tf.keras.backend.variable(0.01, name='var_'+ name)


    def call(self, x):
        return x * self.w 

class SumLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(SumLayer, self).__init__()

    def call(self, x1, x2):
        return x1 + x2


mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

class EmbeddingModel(tf.keras.Model):
    def __init__(self):
        super(EmbeddingModel, self).__init__()

        # place w1 on GPU 0 and create the layer
        with tf.device('/gpu:0'):
            self.L1 = ProdLayer('w1')

        # place w2 on GPU 0 and create the layer        
        with tf.device('/gpu:1'):
            self.L2 = ProdLayer('w2')

        # place w3 on both GPU0 and GPU1 using mirrored scope. Can I do this?
        with mirrored_strategy.scope():
            self.w3 = tf.keras.backend.variable(0.01, name='var_w3')

        # may be do this on CPU? But for now let it perform this on GPU0
        with tf.device('/gpu:0'):
            self.L3 = SumLayer()


    def call(self, input_layer):
        # w1 is on GPU0, w2 is on GPU1 and w3 is placed using mirrored scope on both GPUs
        # y1 = w1 * x + w3
        with tf.device('/gpu:0'):
            y1 = self.L1(input_layer) + self.w3

        # y2 = w2 * x + w3
        with tf.device('/gpu:1'):
            y2 = self.L2(input_layer) + self.w3

        # y = y1 + y2 (i.e. w1 * x + w2 * x + 2 * w3) 
        with tf.device('/gpu:0'):
            y_hat = self.L3(y1, y2)

        return y_hat


def myLoss(y, y_hat):
    return (y_hat-y) * (y_hat-y)

with mirrored_strategy.scope():
    model = EmbeddingModel()
    model.compile(tf.optimizers.Adam(lr=0.0001), 
                      loss=myLoss)

train_dataset = np.random.choice(100, size=(1000,))
model.fit(x=train_dataset.astype(np.float32), y=train_dataset.astype(np.float32), 
       batch_size=1000, 
       epochs=10, shuffle=False, verbose=True)