多维张量切片

Question

首先要做的事：我对 TensorFlow 比较陌生。

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

1. 我有 3 个张量 (x,y,z)，形状为 (?,49,3,3,32) [在哪里？是批量大小]
2. 在每个张量上，我计算第 3 轴和第 4 轴的总和 [因此我最终得到 3 个形状为 (?,49,32) 的张量]
3. 通过对上述 3 个张量 (?,49,32) 执行 argmax (A)，我得到一个 (?,49,32) 张量

现在我想使用这个张量从初始x,y,z张量中选择切片，格式如下：

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

我尝试使用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 执行类似操作的方法：

import tensorflow as tf

# Make example data
tf.random.set_seed(0)
b = 10  # Batch size
x = tf.random.uniform((b, 49, 3, 3, 32))
y = tf.random.uniform((b, 49, 3, 3, 32))
z = tf.random.uniform((b, 49, 3, 3, 32))
# Stack tensors together
data = tf.stack([x, y, z], axis=2)
# Put reduction axes last
data_t = tf.transpose(data, (0, 1, 5, 2, 3, 4))
# Reduce
s = tf.reduce_sum(data_t, axis=(4, 5))
# Find largest sums
idx = tf.argmax(s, 3)
# Make gather indices
data_shape = tf.shape(data_t, idx.dtype)
bb, ii, jj = tf.meshgrid(*(tf.range(data_shape[i]) for i in range(3)), indexing='ij')
# Gather result
output_t = tf.gather_nd(data_t, tf.stack([bb, ii, jj, idx], axis=-1))
# Reorder axes
output = tf.transpose(output_t, (0, 1, 3, 4, 2))
print(output.shape)
# TensorShape([10, 49, 3, 3, 32])