如何将来自 tf.nn.top_k 的索引与 tf.gather_nd 一起使用？

Question

我正在尝试使用从 tf.nn.top_k 返回的索引从第二个张量中提取值。

我已经尝试使用 numpy 类型索引，以及直接使用 tf.gather_nd，但我注意到索引是错误的。

#  temp_attention_weights of shape [I, B, 1, J]
top_values, top_indices = tf.nn.top_k(temp_attention_weights, k=top_k)

# top_indices of shape [I, B, 1, top_k], base_encoder_transformed of shape [I, B, 1, J]

# I now want to extract from base_encoder_transformed top_indices
base_encoder_transformed = tf.gather_nd(base_encoder_transformed, indices=top_indices)  

# base_encoder_transformed should be of shape [I, B, 1, top_k]

我注意到 top_indices 的格式错误，但我似乎无法将其转换为在 tf.gather_nd 中使用，其中最内层的维度用于索引 base_encoder_transformed 中的每个相应元素。有人知道将 top_indices 转换为正确格式的方法吗？

Answer 1

top_indices将仅在最后一个轴上建立索引，您还需要为其余轴添加索引。这很容易tf.meshgrid：

import tensorflow as tf

# Example input data
I = 4
B = 3
J = 5
top_k = 2
x = tf.reshape(tf.range(I * B * J), (I, B, 1, J)) % 7
# Top K
top_values, top_indices = tf.nn.top_k(x, k=top_k)
# Make indices for the rest of axes
ii, jj, kk, _ = tf.meshgrid(
    tf.range(I),
    tf.range(B),
    tf.range(1),
    tf.range(top_k),
    indexing='ij')
# Stack complete index
index = tf.stack([ii, jj, kk, top_indices], axis=-1)
# Get the same values again
top_values_2 = tf.gather_nd(x, index)
# Test
with tf.Session() as sess:
    v1, v2 = sess.run([top_values, top_values_2])
    print((v1 == v2).all())
    # True

Answer 2

我没有看到使用的理由tf.gather_nd。有一个更简单、更快（无需使用tf.meshgrid）的解决方案，tf.gather与batch_dims参数一起使用。

import tensorflow as tf

# Example input data
I = 4
B = 3
J = 5
top_k = 2
x = tf.reshape(tf.range(I * B * J), (I, B, 1, J)) % 7
# Top K
top_values, top_indices = tf.nn.top_k(x, k=top_k)
#Gather indices along last axis
top_values_2 = tf.gather(x, top_indices, batch_dims = 3)

tf.reduce_all(top_values_2 == top_values).numpy()
#True

请注意，batch_dims在本例中为 3，因为我们要从最后一个轴收集数据，并且 x 的秩为 4。