[英]tf.keras manual device placement
迁移到 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
您可以使用 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)
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.