評價此頁

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 為一維時使用。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,則 offsets 的大小等於包的數量加1。最後一個元素是 input 的大小,或最後一個包(序列)的結束索引位置。

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