张量流中的置换优化

Question

问题设置

我有一批4x4矩阵，包含实际值。

L = tf.placeholder('float32', shape=[None, 4, 4], name='pairwise-loss')

我想找到一批4个排列，这样每个4排列最小化了将排列应用为单热掩模的总和。

def get_best_permutation(L_):
    '''
    Returns a batch of permutations `P` such that `tf.reduce_sum(tf.one_hot(P) * L_, axis=(1,2))`
      is minimized for each batch element.

    Args:
        L_: real valued tensor of shape [None, 4, 4]

    Returns:
        P: tf.int64 tensor of shape [None, 4]

    '''
    raise NotImpelmentedError()

排列大小很小且恒定，通常为3-4。但批量大小预计会非常大。理想情况下，一切都应该在图形内完成，而在GPU上则需要更少的数据传输。

编辑可以安全地假设所有矩阵条目都是正数。

背景

这是为了实现与Permutation Invariant Training类似的东西。

快速而肮脏的解决方案

可以预先计算所有可能的排列，因为它们很小，然后并行应用所有这些排列。最后应用tf.argmin找到最好的一个。但是，我想要一个更有效的解决方案。