python - 如何在GPU设备上运行带权重的tensorflow softmax_cross_entropy？

Question

当我在 GPU 设备上运行以下代码时:

import numpy as np
import tensorflow as tf

with tf.device('gpu'):
    logits_ph = tf.placeholder(tf.float32, [None, 2], name="logits")
    labels_ph = tf.placeholder(tf.float32, [None, 2], name="labels")
    w_ph = tf.placeholder(tf.int32, [None], name="w")
    loss = tf.losses.softmax_cross_entropy(
        onehot_labels=labels_ph,
        logits=logits_ph,
        weights=w_ph
    )

with tf.Session() as sess:
    print(sess.run(
        loss,
        feed_dict={
            labels_ph: np.array([[0.0, 1.0], [1.0, 0.0]]),
            logits_ph: np.array([[0.4, 0.7], [0.5, 1.9]]),
            w_ph: np.array([1.0, 1.0])
        }
    ))

失败并出现错误:

InvalidArgumentError (see above for traceback): Cannot assign a device for operation softmax_cross_entropy_loss/assert_broadcastable/is_valid_shape/has_valid_nonscalar_shape/has_invalid_dims/DenseToDenseSetOperation: Could not satisfy explicit device specification '/device:GPU:*' because no supported kernel for GPU devices is available.
Registered kernels:
  device='CPU'; T in [DT_STRING]
  device='CPU'; T in [DT_UINT16]
  device='CPU'; T in [DT_UINT8]
  device='CPU'; T in [DT_INT64]
  device='CPU'; T in [DT_INT32]
  device='CPU'; T in [DT_INT16]
  device='CPU'; T in [DT_INT8]

但它可以在 CPU 和 GPU 上运行，无需权重。谁能帮我理解为什么它不能在 GPU 上使用权重吗？

Answer 1

事实上，tf.losses.softmax_cross_entropy() 没有 GPU 内核。您可以使用 tf.nn.softmax_cross_entropy_with_logits_v2() 来代替，同时传递加权标签:

import numpy as np
import tensorflow as tf

with tf.device('gpu'):
    logits_ph = tf.placeholder(tf.float32, [None, 2], name="logits")
    labels_ph = tf.placeholder(tf.float32, [None, 2], name="labels")
    w_ph = tf.placeholder(tf.float32, [None], name="w")
    weighted_labels = labels_ph * w_ph[..., tf.newaxis]
    loss = tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=weighted_labels, logits=logits_ph)

with tf.Session() as sess:
    print(weighted_labels.eval({labels_ph:np.array([[0.0, 1.0], [1.0, 0.0]]),
                                w_ph:np.array([2.0, 1.0])}))
    print(loss.eval({logits_ph:np.array([[0.4, 0.7], [0.5, 1.9]]),
                     w_ph:np.array([2.0, 1.0]),
                     labels_ph: np.array([[0.0, 1.0], [1.0, 0.0]])}))
# [[0. 2.]
#  [1. 0.]] <-- first sample was multiplied by 2.0, second sample by 1.0

# [1.1087105 1.6204174] <-- loss for each sample