如何快速将2D NumPy数组（灰度图像）重塑为3D堆叠数组？

问题 #

我有一个形状为(M, N)的2D图像，希望将其重塑为(M//m * N//n, m, n)的形状，即将图像的小补丁堆叠成3D数组。目前我使用两层for循环实现了该功能，代码如下：

import numpy as np
a = np.arange(16).reshape(4,4)
b = np.zeros( (4,2,2))

# a = array([[ 0,  1,  2,  3],
#       [ 4,  5,  6,  7],
#       [ 8,  9, 10, 11],
#       [12, 13, 14, 15]])


row_step = 2
col_step = 2
row_indices = np.arange(4, step = 2)
col_indices = np.arange(4, step = 2)

for i, row_idx in enumerate(row_indices):
    for j, col_idx in enumerate(col_indices):
        b[i*2+j, :,:] = a[row_idx:row_idx+row_step, col_idx:col_idx+col_step]

# b = array([[[ 0.,  1.],
#        [ 4.,  5.]],
#
#       [[ 2.,  3.],
#        [ 6.,  7.]],
#
#       [[ 8.,  9.],
#        [12., 13.]],
#
#       [[10., 11.],
#        [14., 15.]]])

请问有没有更快的实现方法？非常感谢！

更快的实现方法 #

方法1：reshape + transpose 矢量化操作 #

完全避免Python循环，利用numpy内置的数组变形和转置操作，底层是优化后的C实现，效率远超循环：

import numpy as np
a = np.arange(16).reshape(4,4)
m, n = 2, 2

# 1. 拆分成块结构：(M//m, m, N//n, n)
temp = a.reshape(a.shape[0]//m, m, a.shape[1]//n, n)
# 2. 调整维度顺序：(M//m, N//n, m, n)
temp = temp.transpose(0, 2, 1, 3)
# 3. 合并前两个维度得到目标形状
b = temp.reshape(-1, m, n)

print(b)
# 输出结果与循环实现完全一致

方法2：np.lib.stride_tricks.as_strided（零拷贝内存视图） #

如果追求极致性能，可以用as_strided创建原数组的内存视图，无需额外拷贝数据，但必须确保原数组维度能被m、n整除，否则会触发内存越界：

import numpy as np
from numpy.lib.stride_tricks import as_strided

a = np.arange(16).reshape(4,4)
m, n = 2, 2
M, N = a.shape

# 计算各维度的内存步长（单位：字节）
strides = (a.strides[0]*m, a.strides[1]*n, a.strides[0], a.strides[1])
# 定义块结构的形状
shape = (M//m, N//n, m, n)

# 创建内存视图
temp = as_strided(a, shape=shape, strides=strides)
# 合并前两个维度
b = temp.reshape(-1, m, n)

print(b)
# 输出结果一致

方法对比 #

• reshape+transpose：代码可读性高，操作安全，适合绝大多数场景，性能足够优秀。
• as_strided：零拷贝，速度最快，但需要手动计算步长，使用不当易引发内存问题，仅适合对性能要求极高且能保证维度合法性的场景。

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