numpy数组移轴、换轴之差

numpy 与改变数组轴序有关的函数有三：

• numpy.swapaxes
• numpy.moveaxis
• numpy.transpose

当想将某轴移去某位置时，应用 moveaxis。情景如医疗图像中 CT、MRI 的形状是 [H, W, L] 三维的，预处理时想沿某维 slice 之，为代码通用，想将 slice axis 移去 axis 0，然后总是沿 axis 0 slice。

此时一个隐藏的 bug 是：以为三维数组移轴用 swapaxes 也行，其实不然！这是混淆了「轴标（此轴叫数字几）」与「轴位（此轴排第几）」所致。更高维（如四维数组 [n, H, W, L]）时此差别就明显了，但三维时有些迷惑性。

本文对比用 swapaxes 与 moveaxis 操作三、四维数组时的差异。另因 moveaxis 页面显示 New in version 1.11.0.，似乎此版本之后的 numpy 才有 moveaxis，故展示一种用 transpose 实现 moveaxis 的写法。

Code

import numpy as np

def transpose_as_moveaxis(arr, src_axis, dest_axis):
	"""用 transpose 实现 moveaxis"""
    if src_axis == dest_axis:
        return arr
    axes = list(range(arr.ndim))
    if src_axis < 0:
        src_axis = arr.ndim + src_axis
    if dest_axis < 0:
        dest_axis = arr.ndim + dest_axis
    assert src_axis >= 0 and dest_axis >= 0
    _fv = axes.pop(src_axis)
    axes.insert(dest_axis, _fv)
    return np.transpose(arr, axes)

print("3d")
x3 = np.zeros((1, 2, 3))
for f, t in zip(
    [0, 0, 1, 1, 2, 2],
    [1, 2, 0, 2, 0, 1]
):
    print('\t', f, '->', t)
    print(np.swapaxes(x3, f, t).shape)
    print(transpose_as_moveaxis(x3, f, t).shape)
    print(np.moveaxis(x3, f, t).shape)

print("4d")
x4 = np.zeros((1, 2, 3, 4))
print(np.swapaxes(x4, 1, -1).shape)
print(transpose_as_moveaxis(x4, 1, -1).shape)
print(np.moveaxis(x4, 1, -1).shape)

输出：

3d
         0 -> 1
(2, 1, 3)
(2, 1, 3)
(2, 1, 3)
         0 -> 2
(3, 2, 1)
(2, 3, 1)
(2, 3, 1)
         1 -> 0
(2, 1, 3)
(2, 1, 3)
(2, 1, 3)
         1 -> 2
(1, 3, 2)
(1, 3, 2)
(1, 3, 2)
         2 -> 0
(3, 2, 1)
(3, 1, 2)
(3, 1, 2)
         2 -> 1
(1, 3, 2)
(1, 3, 2)
(1, 3, 2)

4d
(1, 4, 3, 2)
(1, 3, 4, 2)
(1, 3, 4, 2)

可见，对 3d 数组，在 0 -> 2、2 -> 0 时 swapaxes 效果都与 moveaxis 不同，4d 时此差别更容易想明白。