EmbeddingBag

class torch.nn.EmbeddingBag(num_embeddings, embedding_dim, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False, _weight=None, include_last_offset=False, padding_idx=None, device=None, dtype=None)

计算嵌入‘袋’（bags）的总和或平均值，而无需实例化中间嵌入。

对于长度固定的‘袋’，没有 per_sample_weights，没有等于 padding_idx 的索引，并且输入是 2D 的情况，此类

• 当 mode="sum" 时，等价于 Embedding 后接 torch.sum(dim=1)，

• 当 mode="mean" 时，等价于 Embedding 后接 torch.mean(dim=1)，

• 当 mode="max" 时，等价于 Embedding 后接 torch.max(dim=1)。

然而，EmbeddingBag 比使用这些操作链条更节省时间和内存。

EmbeddingBag 还支持将 per-sample weights 作为 forward 传递的参数。这会在执行由 mode 指定的加权归约（reduction）之前缩放 Embedding 的输出。如果传递了 per_sample_weights，唯一支持的 mode 是 "sum"，它根据 per_sample_weights 计算加权总和。

参数

• num_embeddings (int) – 词嵌入字典的大小

• embedding_dim (int) – 每个嵌入向量的大小

• max_norm (float, 可选) – 如果给定，范数大于 max_norm 的每个嵌入向量将被重新规范化，使其范数为 max_norm。

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

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

• mode (str, 可选) – "sum", "mean" 或 "max"。指定聚合（reduce）‘袋’的方式。"sum" 计算加权总和，考虑了 per_sample_weights。"mean" 计算‘袋’中值的平均值，"max" 计算每个‘袋’中的最大值。默认值："mean"

• sparse (bool, 可选) – 如果为 True，关于 weight 矩阵的梯度将是稀疏张量。有关稀疏梯度的更多详细信息，请参见注意事项（Notes）。注意：当 mode="max" 时不支持此选项。

• include_last_offset (bool, 可选) – 如果为 True，offsets 将包含一个额外的元素，其中最后一个元素等于 indices 的大小。这与 CSR 格式相匹配。

• padding_idx (int, 可选) – 如果指定，位于 padding_idx 的条目不参与梯度计算；因此，位于 padding_idx 的嵌入向量在训练期间不会更新，即它保持为固定的“填充”。对于新构造的 EmbeddingBag，位于 padding_idx 的嵌入向量默认为全零，但可以更新为其他值以用作填充向量。请注意，位于 padding_idx 的嵌入向量被排除在归约（reduction）之外。

变量

weight (Tensor) – 模块的可学习权重，形状为 (num_embeddings, embedding_dim)，从 $\mathcal{N}(0, 1)$ 初始化。

示例

>>> # an EmbeddingBag module containing 10 tensors of size 3
>>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum')
>>> # a batch of 2 samples of 4 indices each
>>> input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)
>>> offsets = torch.tensor([0, 4], dtype=torch.long)
>>> embedding_sum(input, offsets)
tensor([[-0.8861, -5.4350, -0.0523],
        [ 1.1306, -2.5798, -1.0044]])

>>> # Example with padding_idx
>>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum', padding_idx=2)
>>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9], dtype=torch.long)
>>> offsets = torch.tensor([0, 4], dtype=torch.long)
>>> embedding_sum(input, offsets)
tensor([[ 0.0000,  0.0000,  0.0000],
        [-0.7082,  3.2145, -2.6251]])

>>> # An EmbeddingBag can be loaded from an Embedding like so
>>> embedding = nn.Embedding(10, 3, padding_idx=2)
>>> embedding_sum = nn.EmbeddingBag.from_pretrained(
        embedding.weight,
        padding_idx=embedding.padding_idx,
        mode='sum')

forward(input, offsets=None, per_sample_weights=None)

EmbeddingBag 的前向传播（forward pass）。

参数

• input (Tensor) – 包含嵌入矩阵索引‘袋’的 Tensor。

• offsets (Tensor, 可选) – 仅当 input 为 1D 时使用。offsets 决定了 input 中每个‘袋’（序列）的起始索引位置。

• per_sample_weights (Tensor, 可选) – 一个 float / double 权重的 Tensor，或 None 表示所有权重都应视为 1。如果指定，per_sample_weights 必须与 input 的形状完全相同，并且如果 offsets 非 None，则视为具有相同的 offsets。仅支持 mode='sum'。

返回值

Tensor 输出形状为 (B, embedding_dim)。

返回类型

Tensor

注意

关于 input 和 offsets 的注意事项

• input 和 offsets 必须是相同类型，要么是 int，要么是 long

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

• 如果 input 是形状为 (N) 的 1D Tensor，它将被视为多个‘袋’（序列）的拼接。offsets 必须是包含 input 中每个‘袋’起始索引位置的 1D Tensor。因此，对于形状为 (B) 的 offsets，input 将被视为包含 B 个‘袋’。空‘袋’（即长度为 0 的‘袋’）将返回由零填充的向量。

classmethod from_pretrained(embeddings, freeze=True, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, mode='mean', sparse=False, include_last_offset=False, padding_idx=None)

从给定的二维 FloatTensor 创建 EmbeddingBag 实例。

参数

• embeddings (Tensor) – 包含 EmbeddingBag 权重的 FloatTensor。第一个维度被传递给 EmbeddingBag 作为 ‘num_embeddings’，第二个维度作为 ‘embedding_dim’。

• freeze (bool, 可选) – 如果为 True，则该 Tensor 在学习过程中不会更新。等价于设置 embeddingbag.weight.requires_grad = False。默认值：True

• max_norm (float, 可选) – 参见模块初始化文档。默认值：None

• norm_type (float, 可选) – 参见模块初始化文档。默认值 2。

• scale_grad_by_freq (bool, 可选) – 参见模块初始化文档。默认值 False。

• mode (str, 可选) – 参见模块初始化文档。默认值："mean"

• sparse (bool, 可选) – 参见模块初始化文档。默认值：False。

• include_last_offset (bool, 可选) – 参见模块初始化文档。默认值：False。

• padding_idx (int, 可选) – 参见模块初始化文档。默认值：None。

返回类型

EmbeddingBag

示例

>>> # FloatTensor containing pretrained weights
>>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
>>> embeddingbag = nn.EmbeddingBag.from_pretrained(weight)
>>> # Get embeddings for index 1
>>> input = torch.LongTensor([[1, 0]])
>>> embeddingbag(input)
tensor([[ 2.5000,  3.7000,  4.6500]])