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)

計算嵌入“包”（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 處的嵌入向量被排除在縮減計算之外。

返回型別

Tensor

形狀

• 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]])