多维张量切片

Question

首先，我是TensorFlow的新手。

我正在尝试在tensorflow.keras中实现自定义层，当我尝试实现以下目标时，我遇到了相对困难的时间：

1. 我有3个形状为(x,y,z)的张量(?,49,3,3,32) [其中？是批处理大小]
2. [在每个张量上，我计算第3轴和第4轴上的总和[因此，我得到3个形状为(?,49,32)的张量]
3. 通过在上述3个张量(A)上使用argmax (?,49,32)，我得到了一个(?,49,32)张量

现在，我想使用该张量从初始x,y,z张量中以以下形式选择切片：

• A的最后一个维度中的每个元素对应于选定的张量。（aka：0 = X, 1 = Y, 2 = Z）
• A的最后一个维度的索引对应于我想从张量最后一个维度中提取的切片。

我已经尝试使用tf.gather实现以上目标，但我没有运气。然后，我尝试使用一系列的tf.map_fn，这很丑陋，而且计算量大。

为简化以上操作：假设我们有一个形状为（3,3,3,32）的数组。然后，我尝试实现的numpy等效项是：

import numpy as np
x = np.random.rand(3,3,32)
y = np.random.rand(3,3,32)
z = np.random.rand(3,3,32)
x_sums = np.sum(np.sum(x,axis=0),0);
y_sums = np.sum(np.sum(y,axis=0),0);
z_sums = np.sum(np.sum(z,axis=0),0);
max_sums = np.argmax([x_sums,y_sums,z_sums],0)
A = np.array([x,y,z])
tmp = []
for i in range(0,len(max_sums)):
    tmp.append(A[max_sums[i],:,:,i) 
output = np.transpose(np.stack(tmp))

有什么建议吗？ps：我尝试过tf.gather_nd，但我没有运气

Answer 1

这是使用tf.gather_nd可以执行类似操作的方式：

tf.gather_nd