索引后Numpy数组保存速度大幅变慢的原因问询

对Numpy数组执行索引操作后，保存该数组的速度会显著变慢。以下是最小可复现示例：

import time
import numpy as np

def mre(save_path):
    array = np.zeros((245, 233, 6))

    start = time.time()
    for i in range(1000):
        with open(save_path + '/array1_' + str(i), "wb") as file:
            np.save(file, array)
    end = time.time()
    print(f"No indexing: {end - start}s")

    array2 = array[:,:,[0,1,2,3,4,5]]
    start = time.time()
    for i in range(1000):
        with open(save_path + '/array2_' + str(i), "wb") as file:
            np.save(file, array2)
    end = time.time()
    print(f"With indexing: {end - start}s")
    print("Arrays are equal: ", np.array_equal(array, array2))

运行结果如下：

No indexing: 2.9975574016571045s
With indexing: 10.408239126205444s
Arrays are equal:  True

尽管Numpy判定两个数组内容完全相等，但保存索引后的数组耗时明显更长，请问这是什么原因？

核心原因是索引操作后得到的数组内存布局不连续，而原始数组是内存连续的。

• Numpy默认按行优先（C-order）连续存储数组，用np.zeros创建的原始array内存布局是连续的，可通过array.flags.contiguous验证，结果为True。
• 使用array[:,:,[0,1,2,3,4,5]]这类整数数组索引时，得到的array2不是原始内存的视图，而是重新排列后的数组，其内存布局碎片化，array2.flags.contiguous会返回False。
• np.save保存连续数组时，能直接将整块内存写入文件，速度极快；但保存不连续数组时，Numpy需要先隐式将其转换为连续内存布局（等价于执行np.ascontiguousarray），这个额外的转换步骤在循环多次保存时会累积大量耗时，导致整体速度变慢。

如果要优化保存速度，可以显式将不连续数组转为连续数组后再保存：

array2 = np.ascontiguousarray(array[:,:,[0,1,2,3,4,5]])

处理后，保存速度会和原始数组基本一致。

内容的提问来源于stack exchange，提问作者cozeybozey