torch.nn.functional.embedding_bag¶
- torch.nn.functional.embedding_bag(input, weight, offsets=None, max_norm=None, norm_type=2, scale_grad_by_freq=False, mode='mean', sparse=False, per_sample_weights=None, include_last_offset=False, padding_idx=None)[source][source]¶
计算嵌入“包”(bags)的总和、平均值或最大值。
计算过程无需实例化中间嵌入。详见
torch.nn.EmbeddingBag
。注意
在 CUDA 设备上给定张量时,此操作可能会产生非确定性梯度。详见 可复现性 以获取更多信息。
- 参数
input (LongTensor) – 包含嵌入矩阵索引“包”的张量
weight (Tensor) – 嵌入矩阵,其行数等于最大可能索引 + 1,列数等于嵌入维度大小
offsets (LongTensor, optional) – 仅当
input
为 1D 时使用。offsets
决定了input
中每个“包”(序列)的起始索引位置。max_norm (float, optional) – 如果给出,则范数大于
max_norm
的每个嵌入向量将被重新归一化,使其范数为max_norm
。注意:这将原地修改weight
。norm_type (float, optional) –
max_norm
选项中计算 p-范数时的p
值。默认为2
。scale_grad_by_freq (bool, optional) – 如果给出,这将根据 mini-batch 中词语频率的倒数来缩放梯度。默认为
False
。注意:当mode="max"
时,不支持此选项。mode (str, optional) –
"sum"
、"mean"
或"max"
。指定对“包”进行缩减的方式。默认为:"mean"
sparse (bool, optional) – 如果为
True
,则关于weight
的梯度将是稀疏张量。关于稀疏梯度的更多详情,请参阅torch.nn.Embedding
下的注意事项。注意:当mode="max"
时,不支持此选项。per_sample_weights (Tensor, optional) – 一个 float / double 类型的权重张量,或为 None 以表示所有权重均为 1。如果指定,
per_sample_weights
的形状必须与 input 完全相同,并且如果offsets
不为 None,则视其具有相同的offsets
。include_last_offset (bool, optional) – 如果为
True
,则 offsets 的大小等于包的数量 + 1。最后一个元素是输入的大小,或最后一个包(序列)的结束索引位置。padding_idx (int, optional) – 如果指定,
padding_idx
处的项不参与梯度计算;因此,训练期间不会更新padding_idx
处的嵌入向量,即它保持固定的“填充”状态。请注意,padding_idx
处的嵌入向量被排除在缩减计算之外。
- 返回类型
- 形状
input
(LongTensor) 和offsets
(LongTensor, 可选)如果
input
是形状为 (B, N) 的 2D 张量,它将被视为B
个固定长度为N
的“包”(序列),并且将根据mode
的方式返回B
个聚合值。在这种情况下,offsets
将被忽略,且必须为None
。如果
input
是形状为 (N) 的 1D 张量,它将被视为多个“包”(序列)的串联。offsets
必须是包含input
中每个“包”的起始索引位置的 1D 张量。因此,对于形状为 (B) 的offsets
,input
将被视为包含B
个“包”。空包(即长度为 0 的包)返回的向量将填充为零。
weight
(Tensor): 模块的可学习权重,形状为 (num_embeddings, embedding_dim)per_sample_weights
(Tensor, 可选)。与input
具有相同的形状。output
: 聚合后的嵌入值,形状为 (B, embedding_dim)
示例
>>> # an Embedding module containing 10 tensors of size 3 >>> embedding_matrix = torch.rand(10, 3) >>> # a batch of 2 samples of 4 indices each >>> input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9]) >>> offsets = torch.tensor([0, 4]) >>> F.embedding_bag(input, embedding_matrix, offsets) tensor([[ 0.3397, 0.3552, 0.5545], [ 0.5893, 0.4386, 0.5882]]) >>> # example with padding_idx >>> embedding_matrix = torch.rand(10, 3) >>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9]) >>> offsets = torch.tensor([0, 4]) >>> F.embedding_bag(input, embedding_matrix, offsets, padding_idx=2, mode='sum') tensor([[ 0.0000, 0.0000, 0.0000], [-0.7082, 3.2145, -2.6251]])