How convert output tensor to one-hot tensor?

Question

I need to calculate loss from the softmax output vs target. My target is like [0,0,1] and output is [0.3,0.3,0.4] For the purpose, prediction is correct. But a cost function of below type doesn't account for this kind of accuracy

self._output = output = tf.nn.softmax(y)
self._cost = cost = tf.reduce_mean(tf.square( output - tf.reshape(self._targets, [-1])))

How can I easily convert the output [0.3,0.3,0.4] to [0,0,1] in TF itself?

Answer 1

The typical loss function used for comparing two probability distributions is called cross entropy. TensorFlow has the tf.nn.softmax_cross_entropy_with_logits function which implements that loss. In your case, you can simply do :

self._cost = tf.nn.softmax_cross_entropy_with_logits(
                 y, tf.reshape(self._targets, [-1]))

But if you really want to convert [0.3, 0.3, 0.4] to a one-hot representation for a different purpose, you can use the tf.one_hot function as follows :

sess = tf.InteractiveSession()
a = tf.constant([0.3, 0.3, 0.4])
one_hot_a = tf.one_hot(tf.nn.top_k(a).indices, tf.shape(a)[0])
print(one_hot_a.eval())
# prints [[ 0.  0.  1.]]