PyTorch：索引二维张量和行索引的二维张量答案

【问题标题】：PyTorch: index 2D tensor with 2D tensor of row indicesPyTorch：索引二维张量和行索引的二维张量
【发布时间】：2019-09-15 09:18:42
【问题描述】：

我有一个形状为(x, n) 的torch 张量a 和另一个形状为(y, n) 的张量(y, n)，其中y <= x。 b 的每一列都包含a 的行索引序列，我希望能够以某种方式用b 索引a，这样我就获得了一个形状为(y, n) 的张量，其中第 i 列包含 a[:, i][b[:, i]]（不太确定这是否是正确的表达方式）。

这是一个示例（其中x = 5、y = 3 和 n = 4）：

import torch

a = torch.Tensor(
    [[0.1, 0.2, 0.3, 0.4],
     [0.6, 0.7, 0.8, 0.9],
     [1.1, 1.2, 1.3, 1.4],
     [1.6, 1.7, 1.8, 1.9],
     [2.1, 2.2, 2.3, 2.4]]
)

b = torch.LongTensor(
    [[0, 3, 1, 2],
     [2, 2, 2, 0],
     [1, 1, 0, 4]]
)

# How do I get from a and b to c
# (so that I can also assign to those elements in a)?

c = torch.Tensor(
    [[0.1, 1.7, 0.8, 1.4],
     [1.1, 1.2, 1.3, 0.4],
     [0.6, 0.7, 0.3, 2.4]]
)

我无法理解这一点。我正在寻找的是一种不会产生张量 c 的方法，但也让我将与 c 形状相同的张量分配给 a 的元素，c 由该元素组成。

【问题讨论】：

标签： python pytorch

【解决方案1】：

我尝试使用index_select，但它仅支持1-dim 数组作为索引。

bt = b.transpose(0, 1)
at = a.transpose(0, 1)
ct = [torch.index_select(at[i], dim=0, index=bt[i]) for i in range(len(at))]
c  = torch.stack(ct).transpose(0, 1)
print(c)
"""
tensor([[0.1000, 1.7000, 0.8000, 1.4000],
        [1.1000, 1.2000, 1.3000, 0.4000],
        [0.6000, 0.7000, 0.3000, 2.4000]])
"""

这可能不是最好的解决方案，但希望这至少对您有所帮助。

【讨论】：