Python/PyTorch中无嵌套循环实现矩阵行邻域最小值累加的高效方法

你可以完全借助PyTorch的内置向量化函数替代嵌套循环，核心思路是利用1D最小池化或窗口展开+最小值计算批量处理滑动窗口的最小值，避免逐元素循环。以下是两种高效实现方案，适配不同的边界处理逻辑：

方案1：基于min_pool1d的简洁实现（边界补边缘值） #

如果你的逻辑允许对边界位置（首列、末列）的窗口进行边缘值填充（即j=0时窗口包含M[i-1][0], M[i-1][0], M[i-1][1]，j=末列时包含M[i-1][W-2], M[i-1][W-1], M[i-1][W-1]），这种实现最简洁高效：

import torch
import torch.nn.functional as F

# 假设输入是numpy矩阵，先转为PyTorch张量（支持CPU/GPU）
M = torch.tensor(your_numpy_matrix, dtype=torch.float32)
# 如果有GPU，移到GPU加速
# M = M.cuda()

# 提取上一行（i-1行）的所有数据，形状为(H-1, W)
prev_rows = M[:-1, :]
# 对每行左右各补1个边缘值，形状变为(H-1, W+2)
padded_prev = F.pad(prev_rows, (1, 1), mode='replicate')
# 调整维度适配min_pool1d的输入要求(N, C, L)
padded_prev = padded_prev.unsqueeze(1)
# 执行1D最小池化，窗口大小3，步长1，得到每行每个位置的窗口最小值
window_min = F.min_pool1d(padded_prev, kernel_size=3, stride=1).squeeze(1)
# 对原矩阵的第1行到最后一行执行加法操作
M[1:, :] += window_min

# 如果需要转回numpy数组
# your_result_matrix = M.cpu().numpy()

方案2：基于unfold的精确边界实现（边界取有效元素） #

如果原循环中边界位置只取有效元素（j=0时取M[i-1][0:2]的最小值，j=末列时取M[i-1][W-2:W]的最小值），可以用unfold展开窗口后手动拼接边界结果：

import torch

M = torch.tensor(your_numpy_matrix, dtype=torch.float32)
# M = M.cuda()  # 可选GPU加速

prev_rows = M[:-1, :]
H_minus_1, W = prev_rows.shape

# 展开中间列的3元素窗口，形状为(H-1, W-2, 3)
middle_windows = prev_rows.unfold(dimension=1, size=3, step=1)
# 计算中间列的窗口最小值
middle_min = middle_windows.min(dim=2)[0]

# 计算首列的窗口最小值（取前两个元素）
left_min = prev_rows[:, :2].min(dim=1)[0].unsqueeze(1)
# 计算末列的窗口最小值（取最后两个元素）
right_min = prev_rows[:, -2:].min(dim=1)[0].unsqueeze(1)

# 拼接得到完整的窗口最小值矩阵，形状为(H-1, W)
window_min = torch.cat([left_min, middle_min, right_min], dim=1)
# 执行加法操作
M[1:, :] += window_min

性能优势 #

这两种方案都是完全向量化的操作，PyTorch会自动利用CPU的SIMD指令或GPU的并行计算能力，相比双层嵌套循环，速度提升可达数十到上百倍，尤其当矩阵规模较大时效果更明显。

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