Numpy生成三个网格序列操作过慢，求快速优化方案（支持GPU）

我明白你现在的困扰——用循环拼接生成网格序列耗时极长，这是因为没有利用Numpy的向量化特性，反而用了Python循环这种低效的方式。下面给你几个高效的优化方案：

问题根源分析 #

你的原始代码在Python层面循环1000次，每次生成一个1e6元素的全值数组，再反复调用concatenate拼接。这种方式有两个致命问题：

• Python循环本身有不小的开销，1000次循环累加起来非常可观；
• 每次concatenate都需要重新分配内存并复制数据，内存操作的开销随着循环次数剧增，这是导致你耗时5分钟的核心原因。

优化方案1：用Numpy原生网格生成函数（最快最简洁） #

Numpy内置了indices函数，可以直接生成多维网格的索引矩阵，然后展平成一维序列即可，全程无Python循环，完全是C层面的向量化操作：

mom_grid_number = 1000
import numpy as np
import time

ts = time.time()
# 生成三维网格的索引矩阵
i, j, k = np.indices((mom_grid_number, mom_grid_number, mom_grid_number))
# 展平为一维序列
i_grid_seq = i.ravel()
j_grid_seq = j.ravel()
k_grid_seq = k.ravel()
te = time.time()
print(f"耗时：{te - ts:.4f}秒")

这个方案的耗时应该在0.1秒以内，比你的原始代码快上千倍。

优化方案2：用repeat/tile组合（更节省内存） #

如果担心三维矩阵占用内存（1000^3的int32三维数组约4GB），可以用repeat和tile直接构造一维序列，不需要生成三维数组：

mom_grid_number = 1000
import numpy as np
import time

ts = time.time()
# i序列：每个数字重复1e6次（1000*1000）
i_grid_seq = np.repeat(np.arange(mom_grid_number), mom_grid_number**2)
# j序列：每个数字重复1000次，再整体重复1000次
j_grid_seq = np.tile(np.repeat(np.arange(mom_grid_number), mom_grid_number), mom_grid_number)
# k序列：0-999循环1e6次
k_grid_seq = np.tile(np.arange(mom_grid_number), mom_grid_number**2)
te = time.time()
print(f"耗时：{te - ts:.4f}秒")

这个方案的内存占用只有三个一维数组（约12GB，int32），比三维矩阵省不少，速度同样非常快。

GPU加速方案（利用你的设备优势） #

既然你的设备支持GPU，可以用CuPy（和Numpy API完全兼容的GPU计算库）来进一步加速，尤其是当mom_grid_number更大时，GPU的并行计算优势会更明显：

mom_grid_number = 1000
import cupy as cp
import time

ts = time.time()
# 用CuPy生成网格索引（在GPU上计算）
i, j, k = cp.indices((mom_grid_number, mom_grid_number, mom_grid_number))
# 展平序列（仍在GPU上）
i_grid_seq = i.ravel()
j_grid_seq = j.ravel()
k_grid_seq = k.ravel()
# 如果需要转回CPU的Numpy数组，用.get()
# i_grid_seq_np = i_grid_seq.get()
te = time.time()
print(f"GPU耗时：{te - ts:.4f}秒")

CuPy的版本耗时会比Numpy更快，具体取决于你的GPU性能。

验证结果 #

你可以对比输出的序列是否和原始代码一致，比如取前几个元素或者随机位置的元素，确认正确性后再替换。

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