Pytorch argmax 跨多个维度

Question

我有一个 4D 张量，我想获得最后两个维度的 argmax。torch.argmax只接受整数作为“dim”参数，而不接受元组。

我怎样才能做到这一点？

这就是我的想法，但我不知道如何匹配我的两个“索引”张量的维度。original_array是形状 [1, 512, 37, 59]。

max_vals, indices_r = torch.max(original_array, dim=2)
max_vals, indices_c = torch.max(max_vals, dim=2)
indices = torch.hstack((indices_r, indices_c))

Answer 1

正如其他人提到的，最好展平最后两个维度并应用 argmax

original_array = torch.rand(1, 512, 37, 59)
original_flatten = original_array.view(1, 512, -1)
_, max_ind = original_flatten.max(-1)

..您将获得最大值的线性索引。如果您想要最大值的二维指数，您可以使用列数“展开”指数

# 59 is the number of columns for the (37, 59) part
torch.stack([max_ind // 59, max_ind % 59], -1)

这会给你一个(1, 512, 2)最后 2 个维度包含 2D 坐标的位置。