如何在TensorFlow中交换张量的轴？

Question

我有一个形状为(30, 116, 10)的张量，我想要交换前两个维度，以便我有一个形状为(116, 30, 10)的张量。

我看到 numpy 已经实现了这样的函数 (np.swapaxes)，但我在 tensorflow 中没有找到类似的函数。

你有任何想法吗？

- Alexis Rosuel

Answer 1

tf.transpose提供了与np.swapaxes相同的功能，但更加通用。在您的情况下，您可以执行tf.transpose(orig_tensor, [1, 0, 2])，它等同于np.swapaxes(orig_np_array, 0, 1)。

- keveman

---

如果我不知道输入张量的尺寸，但确定要交换最后两个轴，该怎么办？例如，对于形状为 (2, 3, 4, 5) 的输入变量，如何将其转换为 (2, 3, 5, 4)，同时对于形状为 (3, 4, 5, 6, 7) 的输入变量，同样能够将其转换为 (3, 4, 5, 7, 6)？ - Konstantinos Bairaktaris

@KonstantinosBairaktaris 请看我的回答 - joel

Answer 2

使用tf.einsum可以交换轴，即使输入维度的数量未知。例如：

• tf.einsum("ij...->ji...", input)将交换input的前两个维度;
• tf.einsum("...ij->...ji", input)将交换最后两个维度;
• tf.einsum("aij...->aji...", input)将交换第二个和第三个维度；
• tf.einsum("ijk...->kij...", input)将排列前三个维度；

等等。

- ltskv

---

我听说 tf.einsum 非常慢。 - joel

Answer 3

你可以使用tf.linalg.matrix_transpose仅转置最后两个轴，或者更一般地，通过动态计算前导索引和相对索引来交换任意数量的尾随轴以转置所需的轴。

x = tf.ones([5, 3, 7, 11])
trailing_axes = [-1, -2]

leading = tf.range(tf.rank(x) - len(trailing_axes))   # [0, 1]
trailing = trailing_axes + tf.rank(x)                 # [3, 2]
new_order = tf.concat([leading, trailing], axis=0)    # [0, 1, 3, 2]
res = tf.transpose(x, new_order)
res.shape                                             # [5, 3, 11, 7]

- joel