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 上Embedding ，我將如何處理？ （這只是一個例子，層可能與嵌入無關）

如果解決方案涉及子類化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)