Argmax на тензоре и потолке в Tensorflow

Предположим, у меня есть тензор в Tensorflow, его значения такие:

A = [[0.7, 0.2, 0.1],[0.1, 0.4, 0.5]]

Как я могу изменить этот тензор на следующее:

B = [[1, 0, 0],[0, 0, 1]] 

Другими словами, я хочу просто сохранить максимум и заменить его на 1.
Буду признателен за любую помощь.


person Rouzbeh    schedule 29.06.2017    source источник


Ответы (1)


Я думаю, что вы можете решить это с помощью одной строки:

import tensorflow as tf
import numpy as np

x_data = [[0.7, 0.2, 0.1],[0.1, 0.4, 0.5]]
# I am using hard-coded dimensions for simplicity
x = tf.placeholder(dtype=tf.float32, name="x", shape=(2,3))

session = tf.InteractiveSession()

session.run(tf.one_hot(tf.argmax(x, 1), 3), {x: x_data})

Результат тот, который вы ожидаете:

Out[6]: 
array([[ 1.,  0.,  0.],
       [ 0.,  0.,  1.]], dtype=float32)
person Alessandro Suglia    schedule 29.06.2017
comment
Это было умно! Спасибо - person Rouzbeh; 30.06.2017