如何将NumPy分块展平后的1维数组还原为原始2维数组

现有实现的问题 #

原还原函数使用np.split做数组切分，存在不必要的内存拷贝，且包含多步冗余的维度变换，在大尺寸数组场景下性能损耗明显，可利用NumPy维度变换的视图特性做无拷贝优化。

切分逻辑的对称关系 #

现有切分函数本质是两步维度变换，全程不改变内存中元素的排列顺序，因此还原操作只需要反向执行变换步骤即可，不需要做元素拆分：

• 原始数组形状：(8,12,1)（高度、宽度、通道数）
• 切分第一步：reshape为(2,4,3,4,1)，对应维度为「纵向瓦片数、瓦片高度、横向瓦片数、瓦片宽度、通道数」
• 切分第二步：交换第1、2轴（swapaxes(1,2)），得到瓦片数组形状(2,3,4,4,1)，维度为「纵向瓦片数、横向瓦片数、瓦片高度、瓦片宽度、通道数」
• 最后按C序展平为1维数组，每16个连续元素对应一个4x4瓦片。

优化后的还原实现 #

全程使用reshape和swapaxes两个视图操作，无内存拷贝，代码简洁且执行效率更高：

def untiling(flat_array, original_arr):
    h, w, c = original_arr.shape
    tile_h, tile_w = 4, 4
    # 逆序执行切分阶段的维度变换
    return flat_array.reshape(h//tile_h, w//tile_w, tile_h, tile_w, c)\
                      .swapaxes(1, 2)\
                      .reshape(original_arr.shape)

如果不想依赖原始数组对象，也可以直接传入形状参数实现，灵活性更高：

def untiling(flat_array, original_h, original_w, tile_h=4, tile_w=4, channel=1):
    return flat_array.reshape(original_h//tile_h, original_w//tile_w, tile_h, tile_w, channel)\
                      .swapaxes(1, 2)\
                      .reshape(original_h, original_w, channel)

正确性验证 #

可通过以下代码验证还原结果和原数组完全一致：

import numpy as np

# 原始测试数组
a = np.array([
    [0, 0, 1, 3, 4,  5,  12, 0, 0,  0, 0,  0  ],
    [0, 0, 4, 0, 13, 0,  0,  2, 0,  0, 0,  0  ],
    [1, 2, 3, 4, 5,  6,  7,  8, 0,  0, 0,  0  ],
    [5, 4, 9, 0, 3,  0,  7,  2, 0,  0, 0,  0  ],
    [0, 0, 0, 0, 0,  0,  0,  0, 0,  0, 0,  0  ],
    [0, 0, 0, 0, 1,  0,  5,  7, 5,  0, 1,  0  ],
    [0, 0, 0, 0, 0,  5,  12, 3, 0,  4, 12, 3  ],
    [0, 0, 0, 0, 5,  14, 0,  9, 10, 2, 0,  15 ]
])[..., np.newaxis]  # 补充通道维度匹配流程

# 原切分函数
def tiling(arr):
    b = arr.reshape(arr.shape[0] // 4, 4, arr.shape[1] // 4, 4, 1)
    return b.swapaxes(1, 2)

# 模拟完整流程
flat_arr = tiling(a).flatten('C')
external_return = flat_arr.copy()  # 模拟外部程序返回
restore_arr = untiling(external_return, a)

# 验证结果
print(np.array_equal(restore_arr, a))  # 输出True即为还原正确

只要外部程序返回的一维数组保持C序展平的元素排列顺序不变，即使修改了元素值，该还原函数也可以正常工作，仅需要原始数组的形状信息即可，不依赖原始数值。

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