如何用Numpy将(K,M,N)的3D数组重塑为(n_rows*M,n_cols*N)的2D数组？

已知有一个形状为(K, M, N)的Numpy 3D数组arr，其中每个(M, N)子数组对应一张图像，且已确保K = n_rows * n_cols。需要将其重排为形状(n_rows * M, n_cols * N)的2D数组。此前尝试遍历所有3轴排列组合的方式未成功：

for perm in itertools.permutations([0, 1, 2], 3):
        test = arr.transpose(perm).reshape((n_rows * M, n_cols * N))

但使用Einops可以通过以下代码得到预期结果：

test = ein.rearrange(arr, '(r c) h w -> (r h) (c w)', r=n_rows, c=n_cols)

Numpy实现方案 #

无需依赖第三方库，通过先重塑为4D数组、调整轴顺序、再最终重塑即可实现：

import numpy as np

# 假设已定义n_rows, n_cols，且K = n_rows * n_cols
M, N = arr.shape[1], arr.shape[2]
result = arr.reshape(n_rows, n_cols, M, N).transpose(0, 2, 1, 3).reshape(n_rows * M, n_cols * N)

步骤说明 #

1. 第一步reshape：将原(K, M, N)数组拆分为(n_rows, n_cols, M, N)，把原本的K个图像按行列分组；
2. transpose调整轴：将轴顺序从(行组, 列组, 图像高, 图像宽)改为(行组, 图像高, 列组, 图像宽)，这样同一行组内的图像可以按高度方向拼接；
3. 最后reshape：将4D数组直接拉平为目标2D形状，完成行列方向的图像拼接。

此前尝试3轴排列失败的原因是：仅通过3轴的转置无法实现「分组后按维度拼接」的逻辑，必须通过4D数组的中间状态来完成轴的重组。

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