tf.keras 手动放置设备

Question

迁移到 TF2.0 我正在尝试使用 tf.keras 方法来解决问题。在标准 TF 中，我可以使用 with tf.device(...) 来控制操作的位置。

例如，我可能有一个类似于

的模型


model = tf.keras.Sequential([tf.keras.layers.Input(..),
                             tf.keras.layers.Embedding(...),
                             tf.keras.layers.LSTM(...),
                             ...])

假设我希望网络在 CPU 上运行到 Embedding（包括），然后在 GPU 上运行，我将如何处理？（这只是一个例子，层可能与嵌入无关）

如果解决方案涉及子类化 tf.keras.Model 也可以，我不介意不使用 Sequential

Answer 1

您可以使用 Keras 函数 API:

inputs = tf.keras.layers.Input(..)
with tf.device("/GPU:0"):
    model = tf.keras.layers.Embedding(...)(inputs)
outputs = tf.keras.layers.LSTM(...)(model)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

tf.keras 手动放置设备

tf.keras manual device placement

python

keras

tensorflow

tensorflow2.0