多维张量切片

Question

首先：我对TensorFlow比较陌生。你知道吗

我正在尝试在中实现一个自定义层tensorflow.keras公司当我努力实现以下目标时，我的日子相对不好过：

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

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

• A的最后一个维度中的每个元素对应于所选的张量。 （又名：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))

有什么建议吗？ 附言：我试过tf.gather_nd，但运气不好

Answer 1

这就是如何使用^{}执行类似操作的方法：

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])