我有 2 个张量:
A with shape (batch, sequence, vocab) 和 B 具有形状(批次、序列)。
A = torch.tensor([[[ 1., 2., 3.],
[ 5., 6., 7.]],
[[ 9., 10., 11.],
[13., 14., 15.]]])
B = torch.tensor([[0, 2],
[1, 0]])
我想得到以下内容:
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]]
但是以矢量化的方式。我试过 torch.gather 和其他东西,但我不能让它工作。 有人可以帮我吗?
答案 0 :(得分:1)
>>> 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.]])
您好,您可以使用 torch.gather(A, -1, B.unsqueeze(-1)).squeeze(-1)
。
-1
和 A
之间的第一个 B.unsqueeze(-1)
表示您要沿其选择元素的维度。
B.unsqueeze(-1) 中的第二个 -1
是在 B 上加一个暗度,使两个张量的暗度相同,否则你会得到 RuntimeError: Index tensor must have the same number of dimensions as input tensor
。
最后一个 -1
是将结果从 torch.Size([2, 2, 1])
重塑为 torch.Size([2, 2])