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(...),
                             ...])

假设我想让网络一直保持到Embedding上(包括)在CPU上，然后从那里到GPU上，我该怎么做呢？(这只是一个例子，这些层可能与嵌入无关)

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

Answer 1

可以使用Keras functional：

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)