如何在 TensorFlow 中交换张量轴?

How do I swap tensor's axes in TensorFlow?

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

的张量

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

你有什么想法吗?

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

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

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

等等。

您只能使用 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]