pytorch中具有2个相同形状的张量的Argmax索引答案

【问题标题】：Argmax indexing in pytorch with 2 tensors of equal shapepytorch中具有2个相同形状的张量的Argmax索引
【发布时间】：2020-10-19 11:03:37
【问题描述】：

总结问题

我正在使用 pytorch 中的高维张量，我需要使用来自另一个张量的 argmax 值来索引一个张量。所以我需要用张量x的argmax的结果来索引dim [3,4]的张量y和dim [3,4]。如果张量是：

import torch as T
# Tensor to get argmax from
# expected argmax: [2, 0, 1]
x = T.tensor([[1, 2, 8, 3],
              [6, 3, 3, 5],
              [2, 8, 1, 7]])

# Tensor to index with argmax from preivous
# expected tensor to retrieve [2, 4, 9]
y = T.tensor([[0,  1,  2,  3],
              [4,  5,  6,  7],
              [8,  9, 10, 11]])
# argmax
x_max, x_argmax = T.max(x, dim=1)

我想要一个给定 x 或 x_argmax 的 argmax 索引的操作，在相同的索引 x_argmax 索引中检索张量 y 中的值。

描述你的尝试

这是我尝试过的：

# What I have tried
print(y[x_argmax])
print(y[:, x_argmax])
print(y[..., x_argmax])
print(y[x_argmax.unsqueeze(1)])

我已经阅读了很多关于 numpy 索引、基本索引、高级索引和组合索引的内容。我一直在尝试使用组合索引（因为我想要张量的第一维切片和第二维的索引值）。但是我还没有为这个用例想出一个解决方案。

【问题讨论】：

标签： indexing pytorch tensor

【解决方案1】：

您正在寻找torch.gather:

idx = torch.argmax(x, dim=1, keepdim=true)  # get argmax directly, w/o max
out = torch.gather(y, 1, idx)

结果

tensor([[2],
        [4],
        [9]])

【讨论】：

【解决方案2】：

y[T.arange(3), x_argmax]怎么样？

这对我有用...

说明：你调用T.max(x, dim=1)时带走了维度信息，所以需要显式恢复这些信息。

【讨论】：