使用tensorflow，如何对top_k进行操作，然后用原始张量和修改后的top_k创建一个新的张量

Question

我一直在尝试做一些非常简单的事情，但没有成功。我有一个张量（比如形状为

X

的

(None, 128)

包含一些分数，换句话说，每批有 128 个分数。现在我在这里应用

Y = tf.math.top_k(X, k=a).indices

a

表示最高的

a

分数。为了简单起见，让我们考虑一下，

a = 95

。那么张量

Y

的形状将是

(None, 95)

。到这里就可以了。

现在我原来的

data

张量的形状为

(None, 3969, 128)

。我想对具有 top_k 分数的数据进行一些操作。所以我使用以下方法提取数据：

ti = tf.reshape(Y, [Y.shape[-1], -1])   # Here ti is of shape (95, None)
fs = tf.gather(X, ti[:, 0], axis=-1)    # Here fs is of shape (None, 3969, 95)

然后通过说

Z = fs * 0.7    # Here Z is of shape (None, 3969, 95)

进行我的操作。这也很好。

现在我想创建一个新的张量

F

，首先

F

的形状为

(None, 3969, 128)

，包含所有未更改的数据（得分不属于top_k的数据）和修改后的数据（得分低于top_k的数据） top_k 并已在

Z

中进行了修改），但是这些数据的顺序将与原始数据中的顺序相同，即修改后的数据仍应位于其原始位置。 这就是我被困住的地方。

我对 TensorFlow 比较陌生，所以如果我遗漏了任何简单的或不清楚的内容，我深表歉意。已经坚持了几天了。

谢谢！

Answer 1

一种方法是使用

tf.tensor_scatter_nd_update

。不过，您必须将 topk 函数的索引转换为适合您的数据的索引。为此，您可以结合使用 tf.tile 和 tf.unravel_index 将

X

形状转换为

data

形状。

如果我们假设像您一样的 3D 数据，您可以使用与此类似的内容：

# getting the dimensions of the data tensor
# assuming that the shape of X is (B,N) and the shape of data is (B,D,N)
B, D, N = tf.unstack(tf.shape(data))
topk = tf.math.top_k(X, k=k)
# to get the absolute indices in the original tensor, we need to:
# - tile to get indices from (B,N) to (B,D,N)
# - do index arithmetics to get indices on the flattened tensor
topk_idx_tiled = tf.tile(topk.indices[:,None,:], [1,D,1])
flattened_indices = tf.reshape(tf.reshape(tf.range(B*D)*N,(B,D))[...,None] + topk_idx_tiled, -1)
# unraveling to get indices with batch dimensions so that we have compatibility with scatter_nd
sc_idx = tf.transpose(tf.unravel_index(flattened_indices, tf.shape(data)))
# scattering the updates to update the original data 
updates = tf.reshape(tf.tile(topk.values[:,None,:],[1,D,1]),-1)*0.7
F = tf.tensor_scatter_nd_update(data, sc_idx, updates)