如何在 TensorFlow 中交换张量轴？

Posted 2023-02-23

【中文标题】如何在 TensorFlow 中交换张量轴？

【英文标题】：How do I swap tensor's axes in TensorFlow?

【发布时间】：2016-11-07 19:24:13

【问题描述】：

我有一个形状为(30, 116, 10) 的张量，我想交换前两个维度，这样我就有一个形状为(116, 30, 10) 的张量

我看到 numpy 实现了这样的功能 (np.swapaxes)，我在 tensorflow 中搜索了类似的东西，但一无所获。

你有什么想法吗？

【参考方案1】：

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

【讨论】：

如果我不知道输入张量的尺寸但我确定要交换最后两个轴怎么办？就像，我应该对张量变量做什么，以便形状 (2, 3, 4, 5) 的输入最终会成为 (2, 3, 5, 4) 但同样适用于形状 (3, 4, 5, 6, 7) 的输入（并将其转换为 (3, 4, 5, 7, 6)）跨度>

@KonstantinosBairaktaris 看到我的回答

【参考方案2】：

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

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

等等。

【讨论】：

听说tf.einsum真的很慢

【参考方案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]