Tensorflow embedding_lookup渐变在CPU上注册？

2023年6月23日 343次阅读

我有一些相当于稀疏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

无论我选择何种优化器,都会发生这种情况.我在这里错过了什么吗？

最佳答案 如果你使用

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

你应该得到一个错误.那是因为当前没有在GPU上实现embedding_lookup.

点赞