如何加速从大型NumPy数组中选取数据的操作

嗨Erick，我完全理解你处理这种超大NumPy数组时的速度困扰——你的数组算下来有80多亿个元素，直接用fancy indexing复制确实会慢得让人头疼。下面给你几个实用的优化思路，亲测有效的那种：

• 换用更高效的随机数生成器
  你现在用的np.random.randint是旧版的随机API，NumPy现在推荐使用更高效的np.random.Generator接口，生成索引的速度会快不少，尤其是当样本量很大的时候：

  import numpy as np
  N_total_sig = 600
  # 初始化新的随机生成器
  rng = np.random.default_rng()
  rand_Select_sig = rng.integers(N_total_sig, size=N_total_sig)
  

• 调整数组轴顺序，优化内存访问效率
  NumPy默认是按C-order（行优先）存储数组的，也就是最后一个维度的元素在内存里是连续的。你现在要采样的是第三个维度（axis=2），直接索引的话会频繁跳转到分散的内存块，速度很慢。可以先把要采样的维度移到最前面，采样后再转回去，这样内存访问更连续：

  # 转置数组，把要采样的维度放到第一位
  data_transposed = Original_data.transpose(2, 0, 1, 3)
  # 采样后再转回原轴顺序
  sig = data_transposed[rand_Select_sig, :, :, :].transpose(1, 2, 0, 3)
  

  这个方法能显著提升复制的速度，因为连续内存访问的效率比分散访问高很多。

• 避免不必要的内存开销，用预分配数组
  如果需要多次执行采样操作，可以预先分配好结果数组的空间，再把数据赋值进去，比让NumPy自动创建数组略快一些：

  sig = np.empty_like(Original_data)
  sig[:] = Original_data[:, :, rand_Select_sig, :]
  

• 内存不足？试试内存映射（memmap）
  如果你的数组已经大到超出可用内存，系统会自动用磁盘交换空间（swap），这会让速度暴跌。这时候可以用np.memmap把数组存在磁盘上，操作时只加载需要的部分：

  # 假设你的数组是从npy文件加载的，直接用memmap读取
  Original_data = np.memmap('original_data.npy', dtype=np.float32, mode='r', shape=(512,512,600,51))
  # 采样操作和之前一样，但内存占用会小很多
  sig = Original_data[:, :, rand_Select_sig, :]
  # 如果要保存结果，也可以用memmap
  sig_memmap = np.memmap('sampled_data.npy', dtype=np.float32, mode='w+', shape=(512,512,600,51))
  sig_memmap[:] = sig
  

• 有GPU的话，用CuPy起飞
  如果你的机器有NVIDIA GPU，直接用CuPy替代NumPy，GPU的并行计算能力处理这种大规模数组操作速度会快几个数量级：

  import cupy as cp
  # 把数组转到GPU内存
  original_gpu = cp.asarray(Original_data)
  # 用CuPy的随机生成器生成索引
  rng = cp.random.default_rng()
  rand_idx = rng.integers(N_total_sig, size=N_total_sig)
  # 采样操作
  sig_gpu = original_gpu[:, :, rand_idx, :]
  # 需要的话转回CPU内存
  sig = cp.asnumpy(sig_gpu)
  

先试试前两个方法吧，不需要额外依赖就能看到明显的速度提升。如果内存吃紧就用memmap，有GPU的话直接冲CuPy！

备注：内容来源于stack exchange，提问作者Erickzhou