如何在張量流中對一批二維張量進行排序？

Question

我有張量A形狀爲[#batch, #MaxSequence, #Features]，其中第二維的實際長度（可能小於#MaxSequence）存儲在張量L中。我想按每個批次中序列的第三維中的第二個特徵對A進行排序。我看到這個post使用tf.gather和tf.nn.top_k排序2D張量，但我不知道應用在3D案例。我需要使用循環來完成嗎？

Answer 1

我有一些工作，但它可能存在更好的解決方案。我想我的代碼可能是過於複雜的這個簡單的問題。

這個想法是，我們必須將tf.nn.top_k(a[:,:,1].indices（按第三維中的第二個特徵順序）的返回索引轉換爲tf.gather_nd可以使用的值。特別是對於我下面的例子中，我們需要一個轉換張的

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

到

[[[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.stacktf.stack然後我們將結果重塑爲期望的結果。

完整TF代碼如下（get_shape方法被定義here）：在輸出

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

注意，我們有三個例子。在每個例子中，序列由第二個特徵遞減排序。