python - 在前向传递期间对层使用单独的设备

Question

我正在使用一个tf.custom_ops名为PyroNN。我将此运算符包装在 atf.keras.layers.Layer中，以便在使用功能 API 创建的模型中使用它。

问题是，这个算子不使用 tensorflows gpu 内存管理并且消耗大量内存。这会导致错误GPUassert: Out of memory和终止训练。

是否可以将层的计算委托给单独的设备？我会想到这样的事情：

import tensorflow.keras as K
def crossd_unet_distributed():
    inputs = K.layers.Input((200, 488, 488), dtype=tf.float32, name='line_integrals')

    with tf.device('/GPU:0'):
        x = K.layers.Conv3D(8, 3, padding="same", activation='relu')(inputs)

    with tf.device('/GPU:1'):
        # this custom layer messes with gpu memory, so we place it on its own gpu
        x = SpinBackProjector(detector_shape=(488, 488), volume_shape=(256, 256, 256))(x)

    with tf.device('/GPU:0'):
        outputs = K.layers.Conv3D(1, 3, padding="same", activation='relu')(x)

    return K.Model(inputs=inputs, outputs=outputs)

Keras 是否能够跟踪多个 GPU 上的梯度？我尝试阅读分布式培训文档，但我不确定这是否是我需要的。

标签： pythondistributed-computingtf.kerasmulti-gpu