多维张量切片

Question

第一件事：我对 TensorFlow 比较陌生。

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

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

现在我想使用此张量从初始 x,y,z 张量中 select 切片，形式如下：

A最后一维的每个元素对应selected Tensor。（又名：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 做类似事情的方法：

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

多维张量切片

Multidimensional Tensor slicing

python

tensorflow

keras-layer

tensorflow2.0