numpy MaxPool场景下不破坏原数组关联修改滑动窗口最大值的方法

实现思路 #

全程基于as_strided生成的原始共享内存视图操作，不做reshape、flatten等会破坏内存关联的拷贝操作，通过索引直接定位每个滑动窗口的最大值位置赋值，修改会自动同步到原数组。

完整实现代码 #

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

original = np.array([[[7,9,19,18],[10,11,20,16]],[[24,5,18,11],[6,10,45,12]]], dtype=np.float64)
# 生成共享内存的滑动窗口视图
view = as_strided(original, shape=(2,1,2,2,2), strides=(64,32*2,8*2,32,8))

# 1. 取最后两个核维度的最大值一维索引
kernel_flat_inds = view.argmax(axis=(-2, -1))
# 2. 把一维索引转成核内的二维坐标
kh_ind, kw_ind = np.unravel_index(kernel_flat_inds, view.shape[-2:])

# 3. 生成前四个维度的网格索引，定位到每个滑动窗口
batch_idx, channel_idx, ph_idx, pw_idx = np.ogrid[:view.shape[0], :view.shape[1], :view.shape[2], :view.shape[3]]

# 4. 直接对视图对应位置赋值，修改同步到原数组
view[batch_idx, channel_idx, ph_idx, pw_idx, kh_ind, kw_ind] = 1000

# 验证原数组已同步修改
print(original)

实现说明 #

• 所有操作均基于as_strided生成的原始视图，未产生数组拷贝，内存共享关系全程保留
• 用np.unravel_index将argmax返回的一维索引转换为核内的二维坐标，配合前四个维度的网格索引可精准定位每个滑动窗口的最大值位置
• 对视图的赋值操作直接写入共享内存，原数组original会同步更新，完全适配反向传播场景的需求

内容的提问来源于stack exchange，提问作者Sam-gege