如何加速大型3D NumPy数组嵌套循环的处理效率？

问题描述 #

我创建了一个名为tr_mat的大型3D NumPy数组，形状为：

tr_mat.shape
(1024, 536, 21073)

数组说明：该数组包含1024行、536列的数值，共21073个2D矩阵，每个2D矩阵对应一张图像的灰度像素值。tr_mat已完成转置，目标是提取所有2D数组中相同像素位置的时间序列，并将生成的548864个时间序列分别保存为.1D文本文件。

当前使用的代码运行效率极低：处理约120行（共536行）耗时20分钟，且循环速度随迭代增加逐渐变慢。原代码如下：

tr_mat = frame_mat.transpose() # 转置后的3D NumPy数组
# 保存
rangeh = range(0, 1024)
for row, row_n_l in zip(tr_mat, rangeh): # row = 2D图像的像素行
        for ts_pixel, row_n in zip(row, rangeh): # ts_pixel = 3D数组中像素的时间序列（跨单个2D数组）
        # 保存
        with open(f"/volumes/.../TS_Row{row_n_l}_Pixel{row_n}.1D", "w") as file:
                for i in ts_pixel: file.write(f"{i}\n") # 逐行保存时间序列值

可复现代码（需修改输出目录）：

import numpy as np
tr_mat = np.random.random((1024, 536, 21073))

rangeh = range(0, 1024)
for row, row_n_l in zip(tr_mat, rangeh):
    for ts_pixel, row_n in zip(row, rangeh):
        # 保存
        with open("/volumes/../TS_Row{row_n_l}_Pixel{row_n}.1D", "w") as file: # 请调整输出目录
            for i in ts_pixel: file.write(f"{i}\n")

性能瓶颈分析 #

1. 多层Python循环+逐元素写入：Python循环本身效率低下，逐行写入文件会产生大量IO系统调用，随着文件数量增加，文件系统的读写压力持续增大，导致速度越来越慢。
2. 数组遍历逻辑错误：原代码内层循环使用rangeh = range(0,1024)，但tr_mat的第二维度是536（列数），这会导致循环冗余，且变量命名与实际索引含义不匹配，容易引发逻辑问题。
3. 未利用NumPy向量化优势：纯Python遍历NumPy数组元素，完全浪费了NumPy的高效向量化处理能力。

优化方案 #

方案1：用NumPy内置函数替代循环写入 #

NumPy的np.savetxt可以直接将一维数组写入文本文件，无需逐元素循环，大幅提升IO效率。同时修正索引逻辑错误：

import numpy as np
tr_mat = np.random.random((1024, 536, 21073))
output_dir = "/volumes/../" # 替换为实际输出路径

# 遍历所有像素位置
for row_idx in range(tr_mat.shape[0]):
    for col_idx in range(tr_mat.shape[1]):
        # 获取当前像素的时间序列
        pixel_ts = tr_mat[row_idx, col_idx, :]
        # 用savetxt批量写入，每行一个数值
        np.savetxt(f"{output_dir}TS_Row{row_idx}_Pixel{col_idx}.1D", pixel_ts, fmt="%.6f")

优势：用NumPy的C级实现替代Python循环，IO操作批量执行，效率提升数倍到数十倍。

方案2：多进程并行处理 #

由于每个像素的文件写入是独立任务，可利用多进程并行处理，充分发挥CPU多核能力，进一步缩短总耗时：

import numpy as np
from multiprocessing import Pool

tr_mat = np.random.random((1024, 536, 21073))
output_dir = "/volumes/../" # 替换为实际输出路径

def save_pixel_time_series(args):
    row_idx, col_idx = args
    pixel_ts = tr_mat[row_idx, col_idx, :]
    np.savetxt(f"{output_dir}TS_Row{row_idx}_Pixel{col_idx}.1D", pixel_ts, fmt="%.6f")

if __name__ == "__main__":
    # 生成所有像素的索引对
    pixel_indices = [(i, j) for i in range(tr_mat.shape[0]) for j in range(tr_mat.shape[1])]
    # 进程数根据CPU核心数调整，例如8核CPU设为8
    with Pool(processes=8) as pool:
        pool.map(save_pixel_time_series, pixel_indices)

注意：Windows系统下必须将主逻辑放在if __name__ == "__main__":块中，避免进程创建异常。

额外优化建议 #

• 存储介质升级：如果使用机械硬盘（HDD），大量小文件读写会非常缓慢，换成固态硬盘（SSD）可显著提升IO速度。
• 精度调整：如果不需要高精度数值，可修改fmt参数（如%.4f），减少文件大小和写入时间。
• 路径预定义：提前拼接好输出目录路径，避免在循环中重复执行字符串拼接操作。

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