如何拼接张量内部矩阵？实现n×d×h×w×p×p到n×d×ph×pw的转换

我们需要将形状为$n\times d\times h\times w\times p\times p$的张量，把其中的$p\times p$子矩阵按网格方式拼接，最终得到形状为$n\times d\times hp\times wp$的张量（h行的p×p块上下堆叠，w列的块左右拼接）。

输入输出示例 #

输入张量 #

import numpy as np
# 对应形状(1,1,2,2,3,3)的输入
input_arr = np.array([[[[[[ 0,  1,  2],
           [ 3,  4,  5],
           [ 6,  7,  8]],

          [[ 9, 10, 11],
           [12, 13, 14],
           [15, 16, 17]]],


         [[[18, 19, 20],
           [21, 22, 23],
           [24, 25, 26]],

          [[27, 28, 29],
           [30, 31, 32],
           [33, 34, 35]]]]]])

目标输出 #

output_arr = np.array([[[[0,   1,  2,  9, 10, 11],
         [3,   4,  5, 12, 13, 14],
         [6,   7,  8, 15, 16, 17],
         [18, 19, 20, 27, 28, 29],
         [21, 22, 23, 30, 31, 32],
         [24, 25, 26, 33, 34, 35]]]])

当前循环实现 #

已通过循环完成需求，但希望找到无循环的高效方案，循环代码如下：

p = 3
arr = np.arange(1*1*2*2*p*p).reshape(1, 1, 2, 2, p, p)
answer = np.zeros(shape=(1, 1, 2*p, 2*p))

for (n, d, h, w) in np.ndindex(*arr.shape[:4]):
    answer[n, d, h:h+p, w:w+p] = arr[n, d, h, w]

无循环高效解决方案 #

通过**轴变换（transpose）+ 重塑（reshape）**的组合即可实现，无需循环：

import numpy as np

p = 3
arr = np.arange(1*1*2*2*p*p).reshape(1, 1, 2, 2, p, p)

# 调整轴顺序：将原轴(n,d,h,w,p,p)转为(n,d,h,p,w,p)，再重塑为目标形状
result = arr.transpose(0, 1, 2, 4, 3, 5).reshape(arr.shape[0], arr.shape[1], arr.shape[2]*p, arr.shape[3]*p)

# 验证输出是否符合预期
print(np.array_equal(result, output_arr))  # 输出True

原理说明 #

原张量的维度顺序是(n, d, h, w, p, p)，其中：

• h代表p×p块的行数，w代表块的列数
• 最后两个p分别是每个块的行、列维度

通过transpose(0,1,2,4,3,5)将轴调整为(n, d, h, p, w, p)，此时张量的元素排列顺序已经和拼接后的网格顺序一致：先按h方向堆叠每个块的行，再按w方向拼接每个块的列。最后直接用reshape合并维度，就能得到目标形状的张量。

这种方法完全利用numpy的向量化操作，效率远高于循环，适合处理大规模张量。

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