Mam 2 testy:

A z kształtem (partia, sekwencja, VOCAB) i B z kształtem (partia, sekwencja).

A = torch.tensor([[[ 1.,  2.,  3.],
     [ 5.,  6.,  7.]],

    [[ 9., 10., 11.],
     [13., 14., 15.]]])

B = torch.tensor([[0, 2],
    [1, 0]])

Chcę uzyskać następujące czynności:

C = torch.zeros_like(B)
for i in range(B.shape[0]):
   for j in range(B.shape[1]):
      C[i,j] = A[i,j,B[i,j]]

Ale w widoczny sposób. Próbowałem torch.gather i innych rzeczy, ale nie mogę tego sprawy. Czy ktoś może mi pomóc?

1
lalaland 12 marzec 2021, 20:26

1 odpowiedź

Najlepsza odpowiedź
>>> import torch
>>> A = torch.tensor([[[ 1.,  2.,  3.],
...      [ 5.,  6.,  7.]],
... 
...     [[ 9., 10., 11.],
...      [13., 14., 15.]]])
>>> B = torch.tensor([[0, 2],
...     [1, 0]])
>>> A.shape
torch.Size([2, 2, 3])
>>> B.shape
torch.Size([2, 2])
>>> C = torch.zeros_like(B)
>>> for i in range(B.shape[0]):
...    for j in range(B.shape[1]):
...       C[i,j] = A[i,j,B[i,j]]
... 
>>> C
tensor([[ 1,  7],
        [10, 13]])
>>> torch.gather(A, -1, B.unsqueeze(-1))
tensor([[[ 1.],
         [ 7.]],

        [[10.],
         [13.]]])
>>> torch.gather(A, -1, B.unsqueeze(-1)).shape
torch.Size([2, 2, 1])
>>> torch.gather(A, -1, B.unsqueeze(-1)).squeeze(-1)
tensor([[ 1.,  7.],
        [10., 13.]])

Cześć, możesz użyć torch.gather(A, -1, B.unsqueeze(-1)).squeeze(-1). Pierwszy -1 między A i B.unsqueeze(-1) wskazuje wymiar, w którym chcesz wybrać element.

Drugi -1 w B.NSQueeze (-1) jest dodanie jednego przyciemnienia do b, aby wykonać dwa teor te same dims, w przeciwnym razie otrzymasz RuntimeError: Index tensor must have the same number of dimensions as input tensor.

Ostatni -1 ma zmienić rezultat z torch.Size([2, 2, 1]) do torch.Size([2, 2])

1
emily 14 marzec 2021, 13:24