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上的原因是GPU上的tf.nn.embedding_lookup isn't supported：

因此，假设基本的word2vec示例绑定到CPU（＃514），我们可以看到tf.nn.embedding_lookup在GPU上不起作用。因此，内部使用embedding_lookup的操作系统也不支持GPU（例如，nce_loss）。

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