max - 如何有效地检索 Torch 张量中最大值的索引？

Question

假设有一个火炬张量，例如以下形状:

x = torch.rand(20, 1, 120, 120)

我现在想要的是获取每个 120x120 矩阵的最大值的索引。为了简化问题，我先 x.squeeze()使用形状 [20, 120, 120] .然后我想得到火炬张量，它是形状为 [20, 2] 的索引列表.

我怎样才能快速做到这一点？

Answer 1

torch.topk()就是你要找的。从文档中，

torch.topk(input, k, dim=None, largest=True,sorted=True, out=None) -> (Tensor, LongTensor)

返回 k给定的最大元素 input张量沿
给定的维度。

如 dim未给出，则选择输入的最后一个维度。

如 largest是 False然后返回 k 个最小元素。

返回(值，索引)的命名元组，其中索引是原始输入张量中元素的索引。

bool 选项 sorted如果 True , 将确保返回的 k 个元素本身已排序