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。

Answer 1

查看rollaxis，这是一个移动轴的功能，允许您在一个命令中重新排序阵列。如果im的形状为 i，j，k

rollaxis(im, 2)

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

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

new_im_vec = ravel(rollaxis(im, 2))

Answer 2

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

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

当然，如果你想跳过new_im，你可以在一行中完成，而且只有flatten正在进行任何复制。

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

Answer 3

使用einops：

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

优点：