PyTorch批量向量化替代循环调用random_sample结果顺序不一致问题

问题根源 #

你的错误出在reshape的维度展开顺序和循环拼接的维度顺序不匹配，两种实现的拼接逻辑完全不同，才会出现整体统计值一致但逐点误差不为0的情况。

循环实现的拼接逻辑 #

循环遍历views维度，每个view处理后得到(batch, channel, sample_pts, 1)的结果，在第2维（sample_pts维度）拼接，最终第2维的顺序为：

[view0的所有sample_pts, view1的所有sample_pts, ..., viewN-1的所有sample_pts]

向量化实现的逻辑错误 #

你把输入rgb_emb0_tmp从(batch, views, channel, pts, 1)reshape为(batch*views, channel, pts, 1)时，合并后的batch维度顺序为batch0_view0、batch0_view1...batch0_viewN-1、batch1_view0...；
采样后输出维度为(batch*views, channel, sample_pts, 1)，如果直接reshape为(batch, channel, views * sample_pts, 1)，维度展开顺序是先合并sample_pts维度、再合并views维度，最终第2维的顺序为：

[view0的sample_pts0, view1的sample_pts0, ..., viewN-1的sample_pts0, view0的sample_pts1, view1的sample_pts1...]

和循环实现的顺序完全不一致，自然逐点对比有误差。

修复方案 #

采样后先调整维度顺序再reshape即可，修正后的代码如下：

# 向量化预处理
batch, views = rgb_emb0_tmp.shape[0], rgb_emb0_tmp.shape[1]
feat_vec = rgb_emb0_tmp.reshape(batch*views, rgb_emb0_tmp.shape[2], rgb_emb0_tmp.shape[3], 1)
idx_vec = inputs['r2p_ds_nei_idx0'].reshape(batch*views, inputs['r2p_ds_nei_idx0'].shape[2], inputs['r2p_ds_nei_idx0'].shape[3])

# 调用采样函数
vec_out = random_sample(feat_vec, idx_vec) # 维度 [batch*views, channel, sample_pts, 1]

# 调整维度顺序后reshape，和循环结果完全一致
r2p_emb_vec = vec_out.reshape(batch, views, vec_out.shape[1], vec_out.shape[2], 1) \
                      .permute(0, 2, 1, 3, 4) \
                      .reshape(batch, vec_out.shape[1], -1, 1)

修正后再对比abs(r2p_emb_loop - r2p_emb_vec).max()，结果应为0。

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