Question

首先，我是TensorFlow的新手。

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

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

有什么建议吗？ 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])

多维张量切片

1 个答案: