python - numpy数组的重新排序

Question

我想重新排序我的 numpy 数组的维度。以下代码有效，但速度太慢。

for i in range(image_size):
    for j in range(image_size):
        for k in range(3):
            new_im[k, i, j] = im[i, j, k]

在此之后，我对 new_im 进行矢量化：

new_im_vec = new_im.reshape(image_size**2 * 3)

也就是说，我不需要 new_im，我只需要到达 new_im_vec。有一个更好的方法吗？image_size 约为 256。

score 11 · Accepted Answer

查看rollaxis，这是一个移动轴的函数，允许您在单个命令中重新排序数组。如果im形状为i, j, k

rollaxis(im, 2)

应该返回一个形状为k, i, j的数组。

在此之后，您可以展平您的阵列，ravel是为此目的的明确功能。把这一切放在一起，你有一个很好的单线：

new_im_vec = ravel(rollaxis(im, 2))

score 8 · Accepted Answer

new_im = im.swapaxes(0,2).swapaxes(1,2) # First swap i and k, then i and j
new_im_vec = new_im.flatten() # Vectorize

这应该快得多，因为交换轴返回数组的视图，而不是复制元素。

当然，如果你想跳过new_im，你可以在一行中完成，仍然只是flatten在做任何复制。

new_im_vec = im.swapaxes(0,2).swapaxes(1,2).flatten()

score 1 · Accepted Answer

使用 einops：

x = einops.rearrange(x, 'height width color -> color height width')

优点：

3 回答 3