使用TensorFlow反转张量内的两个元素块

Question

假设我有一个大小为[batch，D]的二阶张量（矩阵），并且我想用Dd列元素交换元素的前d列的位置。我如何在TensorFlow中有效地做到这一点？示例：大小为[4,4]的矩阵，我想用剩余的列块交换元素的前d = 2列的位置。输入：1 2 3 41 2 3 41 2 3 41 2 3 4输出：3 4 1 23 4 1 23 4 1 23 4 1 2谢谢。

Answer 1

假设你的矩阵是A那么tf.concat((A[:,d:], A[:,:d]), axis=1)应该给你你想要的东西。我不知道有任何直接的方法可以使它更有效，但是除非它成为您模型中的瓶颈，否则我不会担心。

更新：如果d << D并且您只需要执行一次操作，则由于减少了复制开销（但也会使用更多的内存），因此以下操作可能会更快

# allocate D + d columns

A = tf.Variable(tf.zeros((batch, D + d)))

# set A[:,:D] to initial value

# copy first d columns to the the last d columns

A = A[:,-d:].assign(A[:,:d])

# use slicing to ignore the first d columns

A = A[:,d:]