评价此页

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)[源代码]#

计算 bag(词袋)的嵌入向量的求和、平均值或最大值。

计算过程不会实例化中间的嵌入向量。更多细节请参阅 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) – 如果给定,将通过词语在小批量中的频率的倒数来缩放梯度。默认为 False。注意:当 mode="max" 时,此选项不受支持。

  • mode (str, optional) – "sum", "mean""max"。指定如何聚合词袋。默认值:"mean"

  • sparse (bool, optional) – 如果为 True,则关于 weight 的梯度将是一个稀疏张量。更多关于稀疏梯度的细节请参阅 torch.nn.Embedding 的“Notes”部分。注意:当 mode="max" 时,此选项不支持。

  • per_sample_weights (Tensor, optional) – 一个浮点数/双精度数权重的张量,或者为 None 表示所有权重都取值为 1。如果指定,per_sample_weights 必须与 input 具有完全相同的形状,并且如果 offsets 不为 None,则 per_sample_weights 会被视为具有相同的 offsets

  • include_last_offset (bool, optional) – 如果为 True,则偏移量的数量等于词袋数量 + 1。最后一个元素是输入的数量,或最后一个词袋(序列)的结束索引位置。这与 CSR 格式匹配。在输入为 2D 时将被忽略。默认值:False

  • padding_idx (int, optional) – 如果指定,则 padding_idx 处的条目不计入梯度;因此,padding_idx 处的嵌入向量在训练期间不会被更新,即它保持为一个固定的“pad”。注意,padding_idx 处的嵌入向量不计入聚合。

返回类型:

张量

形状
  • input (LongTensor) 和 offsets (LongTensor, optional)

    • 如果 input 是 2D 的,形状为 (B, N),它将被视为 B 个包(序列),每个包的固定长度为 N,这将根据 mode 返回 B 个聚合值。offsets 被忽略,并且在此情况下要求为 None

    • 如果 input 是 1D 的,形状为 (N),它将被视为多个包(序列)的连接。offsets 要求是包含 input 中每个包起始索引位置的 1D 张量。因此,对于形状为 (B)offsetsinput 将被视为有 B 个包。空包(即长度为 0 的包)将返回填充为零的向量。

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

  • per_sample_weights (Tensor, optional)。形状与 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]])