Tensorflow embedding_lookup 梯度在CPU上注册？

Question

我有一些相当于稀疏softmax的东西:

...
with tf.device('/gpu:0'):
    indices = tf.placeholder(tf.int32, [None, dimsize])
    self._W = weight_variable([self._num_nodes, input_layer_size])
    self._b = bias_variable([self._num_nodes])
    sampled_W = tf.transpose(tf.nn.embedding_lookup(self._W, indices), [0,2,1]) # [batchsize, inputlayersize, dim1size]
    sampled_b = tf.nn.embedding_lookup(self._b, indices) # [batchsize, dim1size]
    ...

但是，当我启用放置日志记录时，我看到多个渐变实例被放置在 CPU 上，例如:

gradients/.../embedding_lookup_1_grad/Size: /job:localhost/replica:0/task:0/cpu:0

I tensorflow/core/common_runtime/simple_placer.cc:819] gradients/.../embedding_lookup_1_grad/Size: /job:localhost/replica:0/task:0/cpu:0

无论我选择什么优化器，都会发生这种情况。我在这里遗漏了什么吗？

Answer 1

如果你使用

tf.Session(config=tf.ConfigProto(allow_soft_placement=False))

你应该得到一个错误。这是因为 embedding_lookup 目前尚未在 GPU 上实现。