按 2d 张量中的值索引 pytorch 4d 张量

【问题标题】：Index pytorch 4d tensor by values in 2d tensor按 2d 张量中的值索引 pytorch 4d 张量
【发布时间】：2018-11-25 20:37:20
【问题描述】：

我有两个 pytorch 张量：

X 与形状 (A, B, C, D)
I 与形状 (A, B)

I 中的值是[0, C) 范围内的整数。

获得形状为(A, B, D) 的张量Y 的最有效方法是什么，例如：

Y[i][j][k] = X[i][j][ I[i][j] ][k]

【问题讨论】：

非常感谢

标签： python pytorch

【解决方案1】：

您可能希望使用torch.gather 进行索引，并使用expand 将I 调整为所需的大小：

eI = I[..., None, None].expand(-1, -1, 1, X.size(3))  # make eI the same for the last dimension
Y = torch.gather(X, dim=2, index=eI).squeeze()

测试代码：

A = 3 
B = 4 
C = 5 
D = 7

X = torch.rand(A, B, C, D)
I = torch.randint(0, C, (A, B), dtype=torch.long)

eI = I[..., None, None].expand(-1, -1, 1, X.size(3))
Y = torch.gather(X, dim=2, index=eI).squeeze()

# manually gather
refY = torch.empty(A, B, D)
for i in range(A):
    for j in range(B):
        refY[i, j, :] = X[i, j, I[i,j], :]

(refY == Y).all()
# Out[]: tensor(1, dtype=torch.uint8)

【讨论】：