Индекс Pytorch 4d тензора по значениям в 2d тензоре

Question

У меня есть два тензора pytorch:

• X с формой (A, B, C, D)
• I с формой (A, B)

Значенияв I - целые числа в диапазоне [0, C).

---

Какой самый эффективный способ получить тензор Y с формой (A, B, D), такой что:

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

Answer 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)