tensorflow - TensorFlow 中的显式 CPU 放置

Question

我发现官方模型示例中有一段代码让我感到困惑。

with tf.device("/cpu:0"):
  embedding = tf.get_variable(
      "embedding", [vocab_size, size], dtype=data_type())
  inputs = tf.nn.embedding_lookup(embedding, input_.input_data)

为什么在这里使用tf.device("/cpu:0")？除了 GPU 内存泄漏的情况，是否还有其他情况需要我们明确指定 CPU 操作？

Answer 1

将嵌入矩阵放在 CPU 上的原因是 tf.nn.embedding_lookup isn't supported在 GPU 上:

So, given the basic word2vec example being bound to CPU (#514), we can see that tf.nn.embedding_lookup doesn't work on GPU. Therefore, ops that use embedding_lookup internally doesn't support GPU either (for example, nce_loss).

这意味着 embedding 变量的 GPU 放置只会导致不必要的数据从主内存传输到 GPU 内存，反之亦然。因此，将变量显式放置在 CPU 上会更有效。