如何高效实现对角核矩阵卷积？附示例代码

填充宽度确认 #

没错，601×601的奇数尺寸核，对称卷积的零填充宽度确实是300。因为要让输入矩阵边缘的元素能和核中心对齐，填充宽度计算公式为(核边长-1)//2，即(601-1)//2=300，这样卷积后输出矩阵的尺寸会和原输入一致。

优化方案 #

你的自定义代码没有充分利用对角核只有主对角线非零的特性，导致计算冗余，而且双重循环+全子矩阵提取的方式效率很低，以下是几个优化方向：

1. 修正基础代码的错误并简化计算 #

先修正你代码里的明显问题，再利用对角核特性减少计算量：

• 去掉print(sub_matrix)：这会极大拖慢运行速度，完全没必要保留
• 修正变量混淆：原代码中d既表示核尺寸又被赋值为对角矩阵，容易出错
• 只计算对角线元素的乘积和：不需要提取整个601×601子矩阵，只取对应位置的对角线元素相乘再求和

优化后的基础版本：

import numpy as np

def convolve_with_diagonal(matrix, diagonal_kernel):
    m, n = matrix.shape
    d = diagonal_kernel.shape[0]
    pad_width = (d - 1) // 2
    # 提取核的对角值（因为是对角核，直接取主对角线即可）
    kernel_diag = np.diag(diagonal_kernel)
    
    padded_matrix = np.pad(matrix, pad_width, mode='constant')
    output = np.zeros_like(matrix, dtype=np.float64)  # 用更高精度避免溢出
    
    # 遍历每个输出位置，只计算对角线元素的乘积和
    for i in range(m):
        for j in range(n):
            # 提取子矩阵的主对角线：从(i,j)到(i+d-1,j+d-1)的对角线元素
            diag_elements = padded_matrix[i:i+d, j:j+d].diagonal()
            output[i, j] = np.sum(diag_elements * kernel_diag)
    
    return output

2. 向量化操作彻底消除双重循环 #

numpy的向量化操作比Python循环快几个数量级，我们可以用as_strided生成所有对角线元素的视图，再进行批量计算：

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

def convolve_with_diagonal_vectorized(matrix, diagonal_kernel):
    m, n = matrix.shape
    d = diagonal_kernel.shape[0]
    pad_width = (d - 1) // 2
    kernel_diag = np.diag(diagonal_kernel)
    
    padded = np.pad(matrix, pad_width, mode='constant')
    # 计算滑动窗口对角线的步长
    stride_row, stride_col = padded.strides
    # 生成所有滑动窗口的对角线元素：形状为(m, n, d)
    diag_windows = as_strided(
        padded,
        shape=(m, n, d),
        strides=(stride_row, stride_col, stride_row + stride_col)
    )
    # 批量计算乘积和
    output = np.sum(diag_windows * kernel_diag, axis=2)
    return output

这个版本完全避免了Python循环，利用numpy的底层C实现加速，速度会比原代码快几十到上百倍。

3. 分块处理减少内存压力 #

70000×70000的矩阵内存占用很高（单精度约18GB，双精度约36GB），可以分块处理，每次只计算一个子块，降低内存占用：

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

def convolve_with_diagonal_blocked(matrix, diagonal_kernel, block_size=2048):
    m, n = matrix.shape
    d = diagonal_kernel.shape[0]
    pad_width = (d - 1) // 2
    kernel_diag = np.diag(diagonal_kernel)
    padded = np.pad(matrix, pad_width, mode='constant')
    output = np.zeros_like(matrix, dtype=np.float64)
    
    # 按块遍历行和列
    for i in range(0, m, block_size):
        for j in range(0, n, block_size):
            # 确定当前块的范围
            i_end = min(i + block_size, m)
            j_end = min(j + block_size, n)
            # 提取对应块的滑动窗口对角线
            padded_block = padded[i:i_end + d - 1, j:j_end + d - 1]
            stride_row, stride_col = padded_block.strides
            diag_windows = as_strided(
                padded_block,
                shape=(i_end - i, j_end - j, d),
                strides=(stride_row, stride_col, stride_row + stride_col)
            )
            output[i:i_end, j:j_end] = np.sum(diag_windows * kernel_diag, axis=2)
    return output

分块大小可以根据你的内存情况调整，比如2048×2048的块，单精度下每个块约16MB，内存压力会小很多。

4. 额外加速建议 #

• 使用numba对循环版本进行JIT编译：如果向量化版本内存还是有压力，可以用numba装饰你的循环代码，能接近C语言的速度
• 利用多线程：numpy默认可能没开多线程，可以设置np.set_num_threads(你的CPU核心数)，或者用numba的并行编译选项
• 降低数据类型：如果精度允许，把矩阵从float64改成float32，能减少一半内存占用

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