PyTorch如何对4D张量的最后两个维度执行argmax操作

原有代码报错原因 #

你的现有代码维度不匹配的核心原因：
第一次对shape为[1, 512, 37, 59]的张量在dim=2（倒数第二个维度）取最大值后，得到的indices_r维度为[1, 512, 59]；第二次对max_vals在dim=2取最大值后，得到的indices_c维度为[1, 512]，二者维度不一致自然无法直接拼接。 #

方案1：拉平最后两维后反解索引（更简洁，推荐） #

这是目前最高效且不易出错的多维度argmax实现方式，适配任意前序维度的张量：

# 1. 拉平最后两个维度，得到shape为 [1, 512, 37*59] 的张量
flatten_last2 = original_array.flatten(start_dim=-2)
# 2. 对拉平后的最后一维取argmax，得到shape为 [1, 512] 的扁平化索引
flat_idx = torch.argmax(flatten_last2, dim=-1)
# 3. 将扁平化索引反解为原始最后两个维度的坐标
indices_r, indices_c = torch.unravel_index(flat_idx, shape=original_array.shape[-2:])
# 4. 可按需拼接为 [1, 512, 2] 的最终索引张量，每个位置对应 (行索引, 列索引)
indices = torch.stack([indices_r, indices_c], dim=-1)

方案2：修正原有分步取max的思路 #

如果你想沿用分步取最大值的逻辑，只需要补一步索引gather操作，拿到对应列位置的行索引即可：

# 第一步：沿倒数第二个维度取最大值，得到每个列位置对应的最大行索引
max_vals_col, indices_r = torch.max(original_array, dim=-2) # 输出shape均为 [1, 512, 59]
# 第二步：沿倒数第一个维度取最大值，得到全局最大的列位置
max_vals_final, indices_c = torch.max(max_vals_col, dim=-1) # 输出shape均为 [1, 512]
# 补一步：用得到的列位置，从indices_r中取出对应的行索引
indices_r_final = indices_r.gather(dim=-1, index=indices_c.unsqueeze(-1)).squeeze(-1)
# 拼接得到最终索引
indices = torch.stack([indices_r_final, indices_c], dim=-1)

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