如何在pytorch中用2-d张量索引3-d张量？答案

【问题标题】：How to index a 3-d tensor with 2-d tensor in pytorch?如何在pytorch中用2-d张量索引3-d张量？
【发布时间】：2020-04-04 16:22:15
【问题描述】：

import torch
a = torch.rand(5,256,120)
min_values, indices = torch.min(a,dim=0)
aa = torch.zeros(256,120)
for i in range(256):
    for j in range(120):
        aa[i,j] = a[indices[i,j],i,j]

print((aa==min_values).sum()==256*120)

我想知道如何避免使用 for-for 循环来获取 aa 值？（我想使用索引来选择另一个 3-d 张量中的元素，所以我不能直接使用 min 返回的值）

【问题讨论】：

标签： pytorch torch

【解决方案1】：

你可以使用 torch.gather

aa = torch.gather(a, 0, indices.unsqueeze(0))

如此处所述：Slicing a 4D tensor with a 3D tensor-index in PyTorch

【讨论】：