如何在tensorflow中对一批2D张量进行排序？

Question

如何在tensorflow中对一批2D张量进行排序？

3

我有一个张量A，形状为[#batch, #MaxSequence, #Features]，其中第二个维度的实际长度（可能小于#MaxSequence）存储在张量L中。我想按每个批次中序列上第三个维度的第二个特征对A进行排序。我看到了这篇帖子，使用tf.gather和tf.nn.top_k对2D张量进行排序，但我不知道如何将其应用于3D情况。我需要使用循环来完成吗？

- Maosi Chen

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- greeness · Accepted Answer

我已经有一些东西在运作，但可能存在更好的解决方案。我猜我的代码可能对于这个简单的问题过于复杂。

思路是，我们需要将tf.nn.top_k(a[:,:,1].indices（按第三维中第二个特征排序）的返回索引转换为tf.gather_nd可以使用的形式。特别是对于下面的例子，我们需要将一个张量转换成：

 [[3, 2, 1, 0],
  [0, 1, 2, 3],
  [2, 0, 3, 1]]

to

 [[[0 3], [0 2], [0 1], [0 0]]
  [[1 0], [1 1], [1 2], [1 3]]
  [[2 2], [2 0], [2 3], [2 1]]]

我发现的内容如下：

- 首先将目标索引展平，得到[3 2 1 0 0 1 2 3 2 0 3 1]。 - 故意构建一个成对的向量[0 0 0 0 1 1 1 1 2 2 2]。 - 使用tf.stack将上述两个向量堆叠起来，然后将结果重塑为所需的形状。

完整的tf代码如下（get_shape方法在此处定义）：

import tensorflow as tf
a = tf.Variable([[[0.51, 0.52, 0.53, 0.94, 0.35],
                  [0.32, 0.72, 0.83, 0.74, 0.55],
                  [0.23, 0.73, 0.63, 0.64, 0.35],
                  [0.01, 0.74, 0.73, 0.04, 0.75]],
                 [[0.32, 0.76, 0.83, 0.74, 0.55],
                  [0.23, 0.72, 0.63, 0.64, 0.35],
                  [0.11, 0.02, 0.03, 0.14, 0.15],
                  [0.01, 0.00, 0.73, 0.04, 0.75]],
                 [[0.51, 0.52, 0.53, 0.94, 0.35],
                  [0.32, 0.00, 0.83, 0.74, 0.55],
                  [0.23, 0.92, 0.63, 0.64, 0.35],
                  [0.11, 0.02, 0.03, 0.14, 0.15]]], tf.float32)
batch_size, seq_length, num_features = get_shape(a)

idx = tf.reshape(range(batch_size), [-1, 1])
idx_flat = tf.reshape(tf.tile(idx, [1, seq_length]), [-1])
top_k_flat = tf.reshape(tf.nn.top_k(a[:,:,1], 
                                    k=seq_length).indices, [-1])
final_idx = tf.reshape(tf.stack([idx_flat, top_k_flat], 1),
                       [batch_size, seq_length, 2])
reordered = tf.gather_nd(a, final_idx)
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
print sess.run(reordered)

                 #ORDERED
                 #by this
                 #Column (within each example)
[[[ 0.01      ,  0.74000001,  0.73000002,  0.04      ,  0.75      ],
  [ 0.23      ,  0.73000002,  0.63      ,  0.63999999,  0.34999999],
  [ 0.31999999,  0.72000003,  0.82999998,  0.74000001,  0.55000001],
  [ 0.50999999,  0.51999998,  0.52999997,  0.94      ,  0.34999999]],
 [[ 0.31999999,  0.75999999,  0.82999998,  0.74000001,  0.55000001],
  [ 0.23      ,  0.72000003,  0.63      ,  0.63999999,  0.34999999],
  [ 0.11      ,  0.02      ,  0.03      ,  0.14      ,  0.15000001],
  [ 0.01      ,  0.        ,  0.73000002,  0.04      ,  0.75      ]],
 [[ 0.23      ,  0.92000002,  0.63      ,  0.63999999,  0.34999999],
  [ 0.50999999,  0.51999998,  0.52999997,  0.94      ,  0.34999999],
  [ 0.11      ,  0.02      ,  0.03      ,  0.14      ,  0.15000001],
  [ 0.31999999,  0.        ,  0.82999998,  0.74000001,  0.55000001]]

请注意，输出结果包含三个示例。在每个示例中，序列按第二个特征值降序排列。