如何用4D核数组高效实现逐元素不同卷积核的2D数组卷积？

需要对n×m的2D信号数组的每个元素，使用对应位置的[i,i]卷积核进行运算，所有卷积核存储在n×m×i×i的4D数组中（[x,y]位置的核对应信号[x,y]元素的运算）。示例如下：

示例 #

2D信号数组 #

[[0 1]
 [5 9]]

4D卷积核数组（每个信号元素对应一个3×3核） #

[[[[1 0 1]  | [[0 0 0]
   [0 1 0]  |  [0 0 0]
   [1 0 1]] |  [0 0 0]]
   ------------------------
  [[1 1 1]  | [[0 1 0]
   [1 1 1]  |  [1 0 1]
   [1 1 1]] |  [0 1 0]]]]

运算过程 #

[[1 0 1]    [[- - -]         [[0 0 0]    [[- - -]
 [0 1 0]  •  [- 0 1]  = 9     [0 0 0]  •  [0 1 -]  = 0
 [1 0 1]]    [- 5 9]]         [0 0 0]]    [5 9 -]]
  
[[1 1 1]    [[- 0 1]          [[0 1 0]    [[0 1 -]
 [1 1 1]  •  [- 5 9]  = 15     [1 0 1]  •  [5 9 -] = 6
 [1 1 1]]    [- - -]]          [0 1 0]]    [- - -]]

期望结果 #

[[9 0]
 [15 6]]

当前低效实现 #

通过循环每个卷积窗口实现，但处理大型数组时速度过慢：

import numpy as np

def fancy_convolve(signal, kernels):
   kernel_size = kernels[0][0].shape[0]
   pad_width = int(kernel_size / 2)
   padded_signal = np.pad(signal, pad_width, 'constant', constant_values=0)
   output = np.empty(signal.shape)
   for x in range(signal.shape[0]):
      for y in range(signal.shape[1]):
         signal_window = padded_signal[x:x+kernel_size, y:y+kernel_size]
         kernel = kernels[x, y]
         output[x, y] = np.dot(signal_window.flatten(), kernel.flatten())
   return output

疑问 #

1. 在Numpy、Scipy或其他库中是否有高效实现该运算的函数？
2. 该运算是否属于卷积范畴？

已查看scipy.ndimage.convolve、scipy.signal.convolve（仅支持单卷积核）及numpy.tensordot（无滑动窗口功能），未找到合适方案。

1. 该运算是否属于卷积范畴？ #

严格来说，这不属于标准卷积。标准卷积是用同一个核在整个信号上滑动，每个位置的输出是核与对应窗口的点积；而你的操作是每个位置使用不同的核与对应窗口做点积，更准确的称呼是「逐位置点积运算」或者「自适应核点积运算」。

2. 高效实现方案 #

方案一：纯Numpy矢量化实现 #

利用numpy.lib.stride_tricks.sliding_window_view生成信号的所有滑动窗口，再和对应位置的核做广播点积后求和，完全避免Python循环：

import numpy as np

def efficient_fancy_convolve(signal, kernels):
    kernel_size = kernels.shape[-1]
    pad_width = kernel_size // 2
    padded_signal = np.pad(signal, pad_width, mode='constant', constant_values=0)
    # 生成所有滑动窗口，形状为(n, m, kernel_size, kernel_size)
    signal_windows = np.lib.stride_tricks.sliding_window_view(padded_signal, (kernel_size, kernel_size))
    # 逐位置点积后求和
    output = np.sum(signal_windows * kernels, axis=(-2, -1))
    return output

该方法速度比原循环快数倍至数十倍，适合处理大型数组。

方案二：Numba加速循环 #

如果环境允许使用Numba，直接对原循环进行JIT编译，可获得接近矢量化的速度，且兼容原有代码逻辑：

import numpy as np
from numba import jit

@jit(nopython=True)
def numba_fancy_convolve(signal, kernels):
    kernel_size = kernels.shape[-1]
    pad_width = kernel_size // 2
    padded_signal = np.pad(signal, pad_width, mode='constant', constant_values=0)
    n, m = signal.shape
    output = np.empty((n, m))
    for x in range(n):
        for y in range(m):
            signal_window = padded_signal[x:x+kernel_size, y:y+kernel_size]
            kernel = kernels[x, y]
            output[x, y] = np.dot(signal_window.flatten(), kernel.flatten())
    return output

首次调用会触发编译，后续调用速度极快。

方案三：Einops简化实现 #

若安装了einops库，可通过更直观的语法实现：

import numpy as np
from einops import reduce

def einops_fancy_convolve(signal, kernels):
    kernel_size = kernels.shape[-1]
    pad_width = kernel_size // 2
    padded_signal = np.pad(signal, pad_width, mode='constant', constant_values=0)
    signal_windows = np.lib.stride_tricks.sliding_window_view(padded_signal, (kernel_size, kernel_size))
    # 逐元素相乘后求和
    output = reduce(signal_windows * kernels, 'n m k1 k2 -> n m', 'sum')
    return output

语法简洁，可读性更强。

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