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)

计算 embedding bags 的总和、均值或最大值。

计算是在不实例化中间 embedding 的情况下完成的。 有关更多详细信息，请参阅 torch.nn.EmbeddingBag。

注意

当在 CUDA 设备上给定张量时，此操作可能会产生非确定性梯度。 有关更多信息，请参阅 可重复性。

参数

• input (LongTensor) – 包含 embedding 矩阵索引 bags 的张量

• weight (Tensor) – embedding 矩阵，行数等于最大可能索引 + 1，列数等于 embedding 大小

• offsets (LongTensor, 可选) – 仅当 input 为 1D 时使用。 offsets 确定 input 中每个 bag（序列）的起始索引位置。

• max_norm (float, 可选) – 如果给定，则将范数大于 max_norm 的每个 embedding 向量重新归一化，使其范数为 max_norm。 注意：这将就地修改 weight。

• norm_type (float, 可选) – 用于计算 max_norm 选项的 p-范数的 p。 默认值为 2。

• scale_grad_by_freq (bool, 可选) – 如果给定，这将按 mini-batch 中单词频率的倒数缩放梯度。 默认值为 False。 注意：当 mode="max" 时，不支持此选项。

• mode (str, 可选) – "sum"、"mean" 或 "max"。 指定减少 bag 的方式。 默认值："mean"

• sparse (bool, 可选) – 如果 True，则相对于 weight 的梯度将是稀疏张量。 有关稀疏梯度的更多详细信息，请参阅 torch.nn.Embedding 下的注释。 注意：当 mode="max" 时，不支持此选项。

• per_sample_weights (Tensor, 可选) – float/double 权重张量，或者 None 表示所有权重应视为 1。 如果指定，则 per_sample_weights 必须与输入具有完全相同的形状，并且被视为具有相同的 offsets（如果这些 offsets 不为 None）。

• include_last_offset (bool, 可选) – 如果 True，则 offsets 的大小等于 bag 的数量 + 1。 最后一个元素是输入的大小，或最后一个 bag（序列）的结束索引位置。

• padding_idx (int, 可选) – 如果指定，则 padding_idx 处的条目不参与梯度计算； 因此，padding_idx 处的 embedding 向量在训练期间不会更新，即它保持为固定的“填充”。 请注意，padding_idx 处的 embedding 向量不包含在归约中。

返回类型

Tensor

形状

• input (LongTensor) 和 offsets (LongTensor, 可选)

  • 如果 input 是形状为 (B, N) 的 2D 张量，它将被视为 B 个 bags（序列），每个 bags 的固定长度为 N，这将返回 B 个值，这些值以取决于 mode 的方式聚合。 在这种情况下，offsets 被忽略，并且必须为 None。

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

• weight (Tensor): 模块的可学习权重，形状为 (num_embeddings, embedding_dim)

• per_sample_weights (Tensor, 可选)。 具有与 input 相同的形状。

• output: 聚合的 embedding 值，形状为 (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]])