在 PyTorch 的 nn.EmbeddingBag 中，offsets 是什么意思？

Question

当偏移量有两个数字时，我知道它的含义，但是当有多于两个数字时，例如：

weight = torch.FloatTensor([[1, 2, 3], [4, 5, 6]])
embedding_sum = nn.EmbeddingBag.from_pretrained(weight, mode='sum')
print(list(embedding_sum.parameters()))
input = torch.LongTensor([0,1])
offsets = torch.LongTensor([0,1,2,1])

print(embedding_sum(input, offsets))

结果是：

[Parameter containing:
tensor([[1., 2., 3.],
        [4., 5., 6.]])]
tensor([[1., 2., 3.],
        [4., 5., 6.],
        [0., 0., 0.],
        [0., 0., 0.]])

who can help me?

- andy

Answer 1

import torch
import torch.nn as nn

weight = torch.FloatTensor([[1, 2, 3], [4, 5, 6]])
embedding_sum = nn.EmbeddingBag.from_pretrained(weight, mode='sum')
print(embedding_sum.weight)

""" output
Parameter containing:
tensor([[1., 2., 3.],
        [4., 5., 6.]])
"""

input = torch.LongTensor([0, 1])
offsets = torch.LongTensor([0, 1, 2, 1])

根据这些偏移量，您将获得以下样本。

"""
sample_1: input[0:1] # tensor([0])
sample_2: input[1:2] # tensor([1])
sample_3: input[2:1] # tensor([])
sample_4: input[1:]  # tensor([1])
"""

将上述示例嵌入

标签中

# tensor([0]) => lookup 0  => embedding_sum.weight[0] => [1., 2., 3.]
# tensor([1]) => lookup 1  => embedding_sum.weight[1] => [4., 5., 6.]
# tensor([])  => empty bag                            => [0., 0., 0.]
# tensor([1]) => lookup 1  => embedding_sum.weight[1] => [4., 5., 6.]

print(embedding_sum(input, offsets))

""" output
tensor([[1., 2., 3.],
        [4., 5., 6.],
        [0., 0., 0.],
        [4., 5., 6.]])
"""

另一个例子：

input = torch.LongTensor([0, 1])
offsets = torch.LongTensor([0, 1, 0])

根据这些偏移量，您将获得以下样本。

"""
sample_1: input[0:1] # tensor([0])
sample_2: input[1:0] # tensor([])
sample_3: input[0:]  # tensor([0, 1])
"""

将上述示例嵌入

标签。

# tensor([0])    => lookup 0 => embedding_sum.weight[0] => [1., 2., 3.]
# tensor([])     => empty bag => [0., 0., 0.]
# tensor([0, 1]) => lookup 0 and 1 then reduce by sum 
#                => embedding_sum.weight[0] + embedding_sum.weight[1] => [5., 7., 9.]

print(embedding_sum(input, offsets))

""" output
tensor([[1., 2., 3.],
        [0., 0., 0.],
        [5., 7., 9.]])
"""

- FelixHo

Answer 2

如源码所示，

return F.embedding(
    input, self.weight, self.padding_idx, self.max_norm,
    self.norm_type, self.scale_grad_by_freq, self.sparse) 

它使用functional embedding bag，其中将offsets参数解释为

offsets（LongTensor，可选）- 仅在输入为1D时使用。offsets确定input中每个包（序列）的起始索引位置。

在EmbeddingBag文档中：

如果输入是形状为(N)的1D张量，则将其视为多个包（序列）的串联。offsets需要是一个1D张量，其中包含input中每个包的起始索引位置。因此，对于形状为(B)的offsets，input将被视为具有B个包。空包（即长度为0的包）将返回由零填充的向量。

最后一句话（"空包（即长度为0的包）将返回由零填充的向量。"）解释了结果张量中的零向量。

- ndrwnaguib