池化嵌入模組

Stable API

class fbgemm_gpu.permute_pooled_embedding_modules.PermutePooledEmbeddings(embs_dims: List[int], permute: List[int], device: device | None = None)

用於沿特徵維度排列嵌入輸出的模組

嵌入輸出張量包含批次中所有特徵的嵌入輸出。它以二維格式表示，其中行是批次大小維度，列是特徵 * 嵌入維度。沿特徵維度排列本質上是沿第二個維度（維度 1）進行排列。

示例

>>> import torch
>>> import fbgemm_gpu
>>> from fbgemm_gpu.permute_pooled_embedding_modules import PermutePooledEmbeddings
>>>
>>> # Suppose batch size = 3 and there are 3 features
>>> batch_size = 3
>>>
>>> # Embedding dimensions for each feature
>>> embs_dims = torch.tensor([4, 4, 8], dtype=torch.int64, device="cuda")
>>>
>>> # Permute list, i.e., move feature 2 to position 0, move feature 0
>>> # to position 1, so on
>>> permute = [2, 0, 1]
>>>
>>> # Instantiate the module
>>> perm = PermutePooledEmbeddings(embs_dims, permute)
>>>
>>> # Generate an example input
>>> pooled_embs = torch.arange(
>>>     embs_dims.sum().item() * batch_size,
>>>     dtype=torch.float32, device="cuda"
>>> ).reshape(batch_size, -1)
>>> print(pooled_embs)
>>>
tensor([[ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.,
         14., 15.],
        [16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27., 28., 29.,
         30., 31.],
        [32., 33., 34., 35., 36., 37., 38., 39., 40., 41., 42., 43., 44., 45.,
         46., 47.]], device='cuda:0')
>>>
>>> # Invoke
>>> perm(pooled_embs)
>>>
tensor([[ 8.,  9., 10., 11., 12., 13., 14., 15.,  0.,  1.,  2.,  3.,  4.,  5.,
          6.,  7.],
        [24., 25., 26., 27., 28., 29., 30., 31., 16., 17., 18., 19., 20., 21.,
         22., 23.],
        [40., 41., 42., 43., 44., 45., 46., 47., 32., 33., 34., 35., 36., 37.,
         38., 39.]], device='cuda:0')

引數:

• embs_dims (List[int]) – 所有特徵的嵌入維度列表。長度 = 特徵數量

• permute (List[int]) – 描述每個特徵如何排列的列表。permute[i] 表示將特徵 permute[i] 排列到位置 i。

• device (Optional[torch.device] = None) – 執行此模組的裝置

__call__(pooled_embs: Tensor) → Tensor

執行沿特徵維度的池化嵌入輸出排列

引數:

pooled_embs (Tensor) – 要排列的嵌入輸出。Shape 為 (B_local, total_global_D)，其中 B_local = 區域性批次大小，total_global_D 是所有特徵（全域性）的總嵌入維度

返回:

排列後的嵌入輸出 (Tensor)。與 pooled_embs 具有相同的 Shape

其他 API