如何在Numpy中高效对所有固定大小子数组执行函数 优化分形压缩速度

你的性能瓶颈确实来自Python层的双层循环，每次循环的numpy小运算叠加Python循环的开销，占了绝大多数耗时。通过Numpy原生向量化操作完全去掉Python层循环，可以获得几十到上百倍的性能提升。

优化后完整代码 #

import numpy as np

def find_block_transform(image, resized_image, x, y, block_size, stride):
    N = block_size * block_size
    # 提前计算输入块的固定值，仅需计算一次
    img = image[y:y+block_size, x:x+block_size]
    R = np.sum(img)
    
    # 生成所有符合步长要求的滑动窗口，一次性提取所有候选域块
    windows = np.lib.stride_tricks.sliding_window_view(resized_image, (block_size, block_size))
    windows = windows[::stride, ::stride]
    
    # 向量化批量计算所有候选块的统计指标
    D = np.sum(windows, axis=(-1, -2))
    DD = np.sum(np.square(windows), axis=(-1, -2))
    DR = np.sum(windows * img[np.newaxis, np.newaxis, :, :], axis=(-1, -2))
    
    # 批量计算所有块的变换参数q、p
    denom_q = DD - np.square(D) / N
    q = (DR - D * R / N) / denom_q
    p = (R - D * q) / N
    
    # 批量计算所有块的MSE
    mse_arr = np.mean(
        np.square(windows * q[..., np.newaxis, np.newaxis] + p[..., np.newaxis, np.newaxis] - img),
        axis=(-1, -2)
    )
    
    # 定位MSE最小的块对应的参数
    min_pos = np.unravel_index(np.argmin(mse_arr), mse_arr.shape)
    best_i = min_pos[0] * stride
    best_j = min_pos[1] * stride
    return BlockTransform(best_i, best_j, q[min_pos], p[min_pos], 0)

注意事项 #

• 该代码依赖numpy 1.20及以上版本提供的sliding_window_view接口，若使用低版本numpy，可替换为as_strided实现滑动窗口：

from numpy.lib.stride_tricks import as_strided
H, W = resized_image.shape
num_i = (H - block_size) // stride + 1
num_j = (W - block_size) // stride + 1
windows = as_strided(
    resized_image,
    shape=(num_i, num_j, block_size, block_size),
    strides=(
        stride * resized_image.strides[0],
        stride * resized_image.strides[1],
        resized_image.strides[0],
        resized_image.strides[1]
    )
)

• 256x256分辨率、block_size=10的场景下，优化后代码耗时通常在几十毫秒级别，相比原始5秒的耗时提升接近百倍。
• 若处理超大图像导致窗口数组内存占用过高，可按行拆分窗口分批计算，相比原始双层循环仍有数量级的性能提升。
• 可额外增加数值保护逻辑，避免denom_q为0时出现除零错误，也可根据业务场景限制q的取值范围（通常分形压缩中q取值在0到1之间）。

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