如何在JAX中高效实现np.lib.stride_tricks.sliding_window_view？

我实现了一种算法，用于计算图像中向量与另一张图像对应像素周围指定窗口内所有向量的皮尔逊相关系数（Pearson correlation coefficient）。纯NumPy版本通过np.lib.stride_tricks.sliding_window_view实现，但大尺寸图像下会产生40GB级别的中间数组，内存开销极大。现有JAX版本性能不及Numba实现，核心瓶颈在于滑动窗口视图的低效实现，同时希望获得NumPy/Numba代码的优化建议。

JAX中无需手动实现sliding_window_view，直接使用**jax.lax.conv_general_dilated_patches**即可——这是XLA底层优化的滑动窗口提取API，比手动dynamic_slice+vmap的方式性能高一个数量级。

改进后的JAX滑动窗口函数 #

from functools import partial
from jax import config
config.update("jax_enable_x64", True)
import jax
import jax.numpy as jnp

@partial(jax.jit, static_argnums=(1,))
def moving_window2d(matrix, window_shape):
    window_h, window_w = window_shape
    # 适配conv_general_dilated_patches的输入格式
    batch_dims = matrix.shape[:-2]
    matrix_reshaped = matrix.reshape(-1, *matrix.shape[-2:], 1)  # (总batch数, H, W, 1)
    # 提取滑动窗口
    patches = jax.lax.conv_general_dilated_patches(
        matrix_reshaped,
        filter_shape=(window_h, window_w),
        window_strides=(1, 1),
        padding="VALID",
        lhs_dilation=(1, 1),
        rhs_dilation=(1, 1),
        dimension_numbers=("NHWC", "HWIO", "NHWC")
    )
    # 恢复原维度结构
    out_shape = (*batch_dims, patches.shape[1], patches.shape[2], window_h, window_w)
    return patches.reshape(out_shape).squeeze(-1)

优化后的JAX版PCC函数 #

@partial(jax.jit, static_argnums=(2, 3))
def pcc_jax_optimized(img1, img2, m, n):
    _n = 2 * n + 1
    _m = 2 * m + 1
    _img2 = img2[..., m:-m, n:-n]
    
    # 预计算中心化向量
    img1_centered = img1 - img1.mean(axis=-3, keepdims=True)
    img2_centered = _img2 - _img2.mean(axis=-3, keepdims=True)
    
    # 提取img1的滑动窗口
    img1_window = moving_window2d(img1_centered, (_m, _n))
    
    # 计算分子：窗口内向量与img2对应向量的点积
    numerator = jnp.sum(img1_window * img2_centered[..., None, None], axis=-5)
    
    # 计算分母：窗口内向量L2范数平方 × img2对应向量L2范数平方的平方根
    img1_window_norm = jnp.sum(img1_window ** 2, axis=-5)
    img2_norm = jnp.sum(img2_centered ** 2, axis=-3)[..., None, None]
    denominator = jnp.sqrt(img1_window_norm * img2_norm)
    
    return numerator / denominator

性能提升原因

• conv_general_dilated_patches直接利用XLA硬件加速，避免手动vmap的调度开销；
• 无需构造滑动窗口起始点数组，减少内存占用和计算步骤；
• 原生支持自动批处理与并行化，适配GPU/TPU等加速设备。

NumPy版本优化：减少内存开销 #

纯NumPy版本的核心问题是sliding_window_view生成的大中间数组，可通过广播+卷积避免显式展开窗口，直接计算所需统计量：

import numpy as np
from scipy.ndimage import uniform_filter

def pcc_numpy_optimized(img1, img2, m, n):
    _n = 2 * n + 1
    _m = 2 * m + 1
    _img2 = img2[..., m:-m, n:-n]
    
    # 中心化处理
    img1_centered = img1 - img1.mean(axis=-3, keepdims=True)
    img2_centered = _img2 - _img2.mean(axis=-3, keepdims=True)
    
    # 利用卷积计算窗口内点积和平方和
    numerator = np.sum([uniform_filter(img1_centered[k] * img2_centered[k], size=(_m, _n), mode='constant') 
                       for k in range(img1.shape[0])], axis=0)
    img1_window_norm = np.sum([uniform_filter(img1_centered[k] ** 2, size=(_m, _n), mode='constant') 
                              for k in range(img1.shape[0])], axis=0)
    img2_norm = np.sum(img2_centered ** 2, axis=-3)
    
    denominator = np.sqrt(img1_window_norm * img2_norm)
    return numerator / denominator

Numba版本优化：减少冗余计算+高效并行 #

现有Numba代码存在冗余遍历和不必要的数组拷贝，优化如下：

import numba as nb

@nb.njit(fastmath=True, parallel=True)
def pcc_numba_optimized(img1, img2, m, n):
    assert img1.shape == img2.shape
    assert img1.ndim == 3
    K, M, N = img1.shape
    _n = 2 * n + 1
    _m = 2 * m + 1
    out_h = M - _m + 1
    out_w = N - _n + 1
    
    # 预计算img1的中心化值和平方值，避免重复计算
    img1_centered = np.empty_like(img1)
    img1_sq = np.empty_like(img1)
    for i in nb.prange(M):
        for j in range(N):
            mean1 = 0.0
            for k in range(K):
                mean1 += img1[k, i, j]
            mean1 /= K
            for k in range(K):
                val = img1[k, i, j] - mean1
                img1_centered[k, i, j] = val
                img1_sq[k, i, j] = val * val
    
    out = np.empty((out_h, out_w, _m, _n), dtype=np.float64)
    
    for i in nb.prange(out_h):
        for j in range(out_w):
            # 计算img2对应点的中心化值和L2范数平方
            mean2 = 0.0
            for k in range(K):
                mean2 += img2[k, i + m, j + n]
            mean2 /= K
            denom2 = 0.0
            img2_centered_k = np.empty(K, dtype=np.float64)
            for k in range(K):
                val = img2[k, i + m, j + n] - mean2
                img2_centered_k[k] = val
                denom2 += val * val
            
            # 遍历窗口内像素计算相关系数
            for a in range(_m):
                for b in range(_n):
                    numerator = 0.0
                    denom1 = 0.0
                    for k in range(K):
                        val1 = img1_centered[k, i + a, j + b]
                        numerator += val1 * img2_centered_k[k]
                        denom1 += img1_sq[k, i + a, j + b]
                    out[i, j, a, b] = numerator / np.sqrt(denom1 * denom2)
    return out

优化点说明

1. 预计算中心化和平方值：避免窗口遍历中重复计算img1的中心化值与平方值；
2. 减少数组拷贝：直接在循环中计算img2的中心化值，无需提前拷贝_img2；
3. 合理并行层级：将prange放在最外层空间维度，最大化并行效率；
4. 合并循环计算：在同一K维度循环中完成均值、中心化、平方和的计算，减少循环次数。

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