вернуть замаскированный softmax top_k каждой строки для 2D-тензора - PullRequest
Новая
       43

вернуть замаскированный softmax top_k каждой строки для 2D-тензора

0 голосов
/ 13 ноября 2018

Для любого 2D-тензора, например

[[2,5,4,7], [7,5,6,8]],

Я хочу сделать softmax для верхнего k элемента в каждой строке, а затем создать новый тензор, заменив все остальные элементы на 0.

Результатом должно быть получение softmax из верхних k (здесь k = 2) элементов для каждой строки [[7,5], [8,7]], что таким образом [[0.880797,0.11920291], [0.7310586,0.26894143]] и затем воссоздать новый тензор по индексу верхних k элементов в исходном тензоре, конечный результат должен быть

[[0,0.11920291,0,0.880797], [0.26894143,0,0,0.7310586]].

Возможно ли реализовать этот вид в маске softmax в тензорном потоке? Большое спасибо заранее!

1 Ответ

0 голосов
/ 13 ноября 2018

Вот как вы можете это сделать: 

import tensorflow as tf

# Input data
a = tf.placeholder(tf.float32, [None, None])
num_top = tf.placeholder(tf.int32, [])
# Find top elements
a_top, a_top_idx = tf.nn.top_k(a, num_top, sorted=False)
# Apply softmax
a_top_sm = tf.nn.softmax(a_top)
# Reconstruct into original shape
a_shape = tf.shape(a)
a_row_idx = tf.tile(tf.range(a_shape[0])[:, tf.newaxis], (1, num_top))
scatter_idx = tf.stack([a_row_idx, a_top_idx], axis=-1)
result = tf.scatter_nd(scatter_idx, a_top_sm, a_shape)
# Test
with tf.Session() as sess:
    result_val = sess.run(result, feed_dict={a: [[2, 5, 4, 7], [7, 5, 6, 8]], num_top: 2})
    print(result_val)

Вывод: 

[[0.         0.11920291 0.         0.880797  ]
 [0.26894143 0.         0.         0.7310586 ]]

РЕДАКТИРОВАТЬ:

На самом деле, есть функция, которая болеевнимательно следит за тем, что вы намереваетесь, tf.sparse.softmax.Тем не менее, он требует SparseTensor в качестве входных данных, и я не уверен, что он должен быть быстрее, поскольку он должен выяснить, какие разреженные значения объединяются в softmax.Хорошая вещь об этой функции в том, что вы можете иметь разное количество элементов для softmax в каждом ряду, но в вашем случае это не кажется важным.Во всяком случае, вот реализация с этим, если вы найдете это полезным. 

import tensorflow as tf

a = tf.placeholder(tf.float32, [None, None])
num_top = tf.placeholder(tf.int32, [])
# Find top elements
a_top, a_top_idx = tf.nn.top_k(a, num_top, sorted=False)
# Flatten values
sparse_values = tf.reshape(a_top, [-1])
# Make sparse indices
shape = tf.cast(tf.shape(a), tf.int64)
a_row_idx = tf.tile(tf.range(shape[0])[:, tf.newaxis], (1, num_top))
sparse_idx = tf.stack([a_row_idx, tf.cast(a_top_idx, tf.int64)], axis=-1)
sparse_idx = tf.reshape(sparse_idx, [-1, 2])
# Make sparse tensor
a_top_sparse = tf.SparseTensor(sparse_idx, sparse_values, shape)
# Reorder sparse tensor
a_top_sparse = tf.sparse.reorder(a_top_sparse)
# Softmax
result_sparse = tf.sparse.softmax(a_top_sparse)
# Convert back to dense (or you can keep working with the sparse tensor)
result = tf.sparse.to_dense(result_sparse)
# Test
with tf.Session() as sess:
    result_val = sess.run(result, feed_dict={a: [[2, 5, 4, 7], [7, 5, 6, 8]], num_top: 2})
    print(result_val)
    # Same as before
...