如何自定义NumPy的reshape过程？解决数组转维后概率分布保留问题

我明白你的困扰——你需要确保每个空间位置(h,w)对应的一对二类概率在重塑和还原过程中不被拆分，始终保持和为1的特性。问题的核心在于正确处理维度的顺序，以及澄清对“概率分布保留”的理解。

正确的重塑流程（从(1,2,H,W)到(2,H*W)） #

你的原始数组结构是(1,2,H,W)，其中第一个维度是冗余的（仅包含1个元素），每个空间位置(h,w)对应的二类概率存储在I[0,0,h,w]和I[0,1,h,w]中（两者和为1）。要将其转为(2,H*W)且保留概率对的对应关系，按以下步骤操作：

1. 挤压冗余的第一个维度：
   先用squeeze去掉仅含1个元素的维度，得到(2,H,W)的数组：

   import numpy as np
   I_squeezed = np.squeeze(I, axis=0)
   

2. 展平空间维度：
   将H和W维度合并为H*W，用reshape并指定-1让NumPy自动计算展平后的长度：

   I_re = I_squeezed.reshape(2, -1)
   

   这样处理后，I_re[:, idx]就对应原始数组中(h,w)位置的一对概率值（idx = h*W + w，行优先展平），每个列的两个元素之和始终为1，完美保留了每个空间位置的概率分布。

处理后还原为原始形状（从(2,H*W)到(1,2,H,W)） #

当你得到处理后的I_re_processed（形状(2,H*W)），按以下步骤还原：

1. 恢复空间维度：
   将H*W维度拆回H和W：

   I_processed_reshaped = I_re_processed.reshape(2, H, W)
   

2. 恢复冗余的第一个维度：
   用expand_dims重新添加第一个维度，回到原始的(1,2,H,W)形状：

   I_final = np.expand_dims(I_processed_reshaped, axis=0)
   

关键误解澄清 #

你提到“重塑后第一axis元素和仍为1”，这里可能存在混淆：我们需要保留的是每个空间位置的一对概率和为1（即每个列的和为1），而非每行的和为1。每行的和是所有空间位置中该类概率的总和，自然不会等于1（除非所有位置的概率都是0.5）。

示例验证 #

用一个小例子直观验证：

H, W = 2, 2
# 构造符合要求的原始数组
I = np.array([[[[0.3, 0.6], [0.8, 0.1]], [[0.7, 0.4], [0.2, 0.9]]]])
# 检查每个空间位置的概率和：0.3+0.7=1，0.6+0.4=1，0.8+0.2=1，0.1+0.9=1

# 重塑
I_squeezed = np.squeeze(I, axis=0)
I_re = I_squeezed.reshape(2, -1)
# I_re结果：
# [[0.3, 0.6, 0.8, 0.1],
#  [0.7, 0.4, 0.2, 0.9]]
# 检查每列和：0.3+0.7=1，0.6+0.4=1，完全符合要求

# 还原
I_processed_reshaped = I_re.reshape(2, H, W)
I_final = np.expand_dims(I_processed_reshaped, axis=0)
# I_final与原始I完全一致

内容的提问来源于stack exchange，提问作者K.Wanter