获取非连续高维张量每个2x2内核最大值索引并同步修改原张量值

核心思路 #

• 仅用reshape生成的副本计算最大值索引，不需要对副本赋值
• 将得到的展平索引转换为最后两个维度的坐标后，直接对原共享内存的视图a赋值，修改会自动同步到original

完整实现代码 #

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)
a = as_strided(original, shape=(2,1,2,2,2), strides=(64, 32*2, 8*2, 32, 8), writeable=True)

# 取内核的宽高参数
kernel_h, kernel_w = a.shape[-2], a.shape[-1]
# reshape仅用于计算索引，不需要共享内存
flatten_for_idx = a.reshape(*a.shape[:-2], kernel_h * kernel_w)
max_flat_idx = flatten_for_idx.argmax(-1)

# 构造前序所有维度的索引
prefix_indices = np.indices(a.shape[:-2])
# 把展平索引转换为最后两维的坐标
max_h_idx = max_flat_idx // kernel_w
max_w_idx = max_flat_idx % kernel_w

# 直接对原视图a赋值，自动同步到original
a[tuple(prefix_indices) + (max_h_idx, max_w_idx)] = 1000

结果验证 #

打印original可以看到对应最大值位置已被修改：

>>> print(original)
[[[   7.    9.   19.   18.]
  [  10. 1000.   20.   16.]]

 [[1000.    5.   18.   11.]
  [   6.   10. 1000.   12.]]

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