如何高效提取NumPy二维矩阵的多个二维切片？

核心思路 #

利用numpy的广播机制，一次性生成所有需要提取的元素的坐标，通过高级索引直接从原矩阵中取出所有裁剪结果，彻底规避Python循环带来的性能开销。

实现代码 #

import numpy as np

# 假设已定义X、from_row、to_row、from_col、to_col
rect_height = to_row[0] - from_row[0]
rect_width = to_col[0] - from_col[0]

# 生成行、列的固定偏移量
row_offsets = np.arange(rect_height)
col_offsets = np.arange(rect_width)

# 计算所有裁剪区域的行/列索引数组
# 广播后得到 (n_rects, rect_height) 和 (n_rects, rect_width) 的索引矩阵
row_indices = from_row[:, None] + row_offsets
col_indices = from_col[:, None] + col_offsets

# 扩展索引形状以支持广播索引，一次性提取所有裁剪结果
crops = X[row_indices[:, :, None], col_indices[:, None, :]]

效率说明 #

原循环方案在Python层面执行500万次迭代，每次迭代的解释器开销会极大拖慢整体速度。而向量化方案直接在numpy的C语言内核中完成所有坐标计算和数据提取，完全绕开了Python循环的性能瓶颈，处理大规模数据时效率提升非常显著。

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