Numpy高维数组多维度应用标量输出函数的高效实现方法

解决方案 #

方案1：优先使用支持指定轴的numpy内置函数 #

如果你用的函数是numpy自带的统计类函数（比如均值、最大值、中位数、方差等），直接指定axis=(-2, -1)即可一步得到结果，性能最高，完全没有Python循环开销：

# 示例：求每个2D图像的均值，输出形状为(x.shape[0], x.shape[1])
result = x.mean(axis=(-2, -1))

所有支持axis参数的numpy函数都可以用这个方法，是最优选择。

方案2：自定义函数用带signature参数的np.vectorize #

普通np.vectorize默认处理标量，但是通过signature参数可以指定输入输出的数组形状，直接处理2D子数组输出标量：

import numpy as np

# 示例自定义函数f，输入2D数组输出标量
def f(img):
    return np.median(img) - img.mean()

# 声明矢量化函数，指定输入为(h,w)的2D数组，输出为标量
f_vec = np.vectorize(f, signature='(h,w)->()')
result = f_vec(x)

输出result的形状就是(x.shape[0], x.shape[1])，写法比嵌套列表推导简洁很多，不需要手动管理循环维度。

方案3：手动展平批量维度再reshape #

如果觉得np.vectorize的性能不够（本质还是Python层循环），可以手动把前两个批量维度展平，单次循环处理所有子数组再还原维度：

# 展平前两个维度，形状变为 (x.shape[0]*x.shape[1], h, w)
x_flat = x.reshape(-1, x.shape[2], x.shape[3])
# 循环处理所有子数组
result_flat = np.array([f(sub_img) for sub_img in x_flat])
# 还原批量维度
result = result_flat.reshape(x.shape[0], x.shape[1])

这个方法适合批量维度更多的场景，比如x是5维数组，前3个维度都是批量维度，只要修改reshape参数即可，不用写多层循环。

性能最优方案 #

如果可以修改自定义函数f，建议把f改造成支持批量输入的形式：即输入形状为(batch, h, w)的3D数组，直接输出形状为(batch,)的数组，全程用numpy矢量化操作实现，完全没有Python循环开销，性能比前面几种方法高几个数量级。

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