torchrec.modules¶

Torchrec 常用模組

torchrec 模組包含各種模組的集合。

這些模組包括

nn.Embedding 和 nn.EmbeddingBag 的擴展，分別稱為 EmbeddingBagCollection 和 EmbeddingCollection。
已建立的模組，例如 DeepFM 和 CrossNet。
常見的模組模式，例如 MLP 和 SwishLayerNorm。
TorchRec 的自訂模組，例如 PositionWeightedModule 和 LazyModuleExtensionMixin。
EmbeddingTower 和 EmbeddingTowerCollection，傳遞給提供的交互模組的邏輯「塔」嵌入。

torchrec.modules.activation¶

激活模組

class torchrec.modules.activation.SwishLayerNorm(input_dims: Union[int, List[int], Size], device: Optional[device] = None)¶

基於： Module

使用層歸一化套用 Swish 函數： Y = X * Sigmoid(LayerNorm(X))。

參數：:

input_dims (Union[int, List[int], torch.Size]) – 要歸一化的維度。如果輸入張量的形狀為 [batch_size, d1, d2, d3]，則設定 input_dim=[d2, d3] 將在最後兩個維度上執行層歸一化。
device (Optional[torch.device]) – 預設計算裝置。

範例

sln = SwishLayerNorm(100)

forward(input: Tensor) → Tensor¶

參數：:: input (torch.Tensor) – 輸入張量。
返回：:: 輸出張量。
返回類型：:: torch.Tensor

training: bool¶

torchrec.modules.crossnet¶

CrossNet API

class torchrec.modules.crossnet.CrossNet(in_features: int, num_layers: int)¶

基於： Module

交叉網路:

交叉網路是在形狀為 \((*, N)\) 的張量上進行一系列「交叉」運算，以得到相同形狀的張量，從而有效地建立了輸入張量上的 \(N\) 個可學習多項式函數。

在這個模組中，交叉運算的定義基於一個滿秩矩陣 (NxN)，這樣交叉效應就可以覆蓋每一層的所有位元。在每一層 l 上，張量被轉換為

\[x_{l+1} = x_0 * (W_l \cdot x_l + b_l) + x_l\]

其中 \(W_l\) 是一個方形矩陣 \((NxN)\)，\(*\) 表示逐元素乘法，\(\cdot\) 表示矩陣乘法。

參數：:

in_features (int) – 輸入的維度。
num_layers (int) – 模組中的層數。

範例

batch_size = 3
num_layers = 2
in_features = 10
input = torch.randn(batch_size, in_features)
dcn = CrossNet(num_layers=num_layers)
output = dcn(input)

forward(input: Tensor) → Tensor¶

參數：:: input (torch.Tensor) – 形狀為 [batch_size, in_features] 的張量。
返回：:: 形狀為 [batch_size, in_features] 的張量。
返回類型：:: torch.Tensor

training: bool¶

class torchrec.modules.crossnet.LowRankCrossNet(in_features: int, num_layers: int, low_rank: int = 1)¶

基於： Module

低秩交叉網路 (Low Rank Cross Net) 是一種高效的交叉網路。它在每一層不是使用全秩交叉矩陣 (NxN)，而是使用兩個核心 \(W (N x r)\) 和 \(V (r x N)\) 來簡化矩陣乘法，其中 r << N。

在每一層 l 上，張量會被轉換為

\[x_{l+1} = x_0 * (W_l \cdot (V_l \cdot x_l) + b_l) + x_l\]

其中 \(W_l\) 是一個向量，\(*\) 表示元素乘法，\(\cdot\) 表示矩陣乘法。

備註

秩 r 的選擇應慎重。通常，我們預期 r < N/2 可以節省計算量；我們預期 \(r ~= N/4\) 可以保持全秩交叉網路的準確性。

參數：:

in_features (int) – 輸入的維度。
num_layers (int) – 模組中的層數。
low_rank (int) – 交叉矩陣的秩設定（預設值 = 1）。值必須始終 >= 1。

範例

batch_size = 3
num_layers = 2
in_features = 10
input = torch.randn(batch_size, in_features)
dcn = LowRankCrossNet(num_layers=num_layers, low_rank=3)
output = dcn(input)

forward(input: Tensor) → Tensor¶

參數：:: input (torch.Tensor) – 形狀為 [batch_size, in_features] 的張量。
返回：:: 形狀為 [batch_size, in_features] 的張量。
返回類型：:: torch.Tensor

training: bool¶

class torchrec.modules.crossnet.LowRankMixtureCrossNet(in_features: int, num_layers: int, num_experts: int = 1, low_rank: int = 1, activation: ~typing.Union[~torch.nn.modules.module.Module, ~typing.Callable[[~torch.Tensor], ~torch.Tensor]] = <built-in method relu of type object>)¶

基於： Module

低秩混合交叉網路 (Low Rank Mixture Cross Net) 是這篇論文中 DCN V2 的實作

LowRankMixtureCrossNet 將每層的可學習交叉參數定義為一個低秩矩陣 \((N*r)\) 以及專家混合 (mixture of experts)。與 LowRankCrossNet 不同的是，它不是依賴單一專家來學習特徵交叉，而是利用 \(K\) 個專家；每個專家都在不同的子空間中學習特徵交互作用，並使用依賴於輸入 \(x\) 的閘控機制自適應地組合學習到的交叉。

在每一層 l 上，張量會被轉換為

\[x_{l+1} = MoE({expert_i : i \in K_{experts}}) + x_l\]

每個 \(expert_i\) 定義為

\[expert_i = x_0 * (U_{li} \cdot g(C_{li} \cdot g(V_{li} \cdot x_l)) + b_l)\]

其中 \(U_{li} (N, r)\)、\(C_{li} (r, r)\) 和 \(V_{li} (r, N)\) 是低秩矩陣，\(*\) 表示元素乘法，\(x\) 表示矩陣乘法，\(g()\) 是非線性激活函數。

當 num_expert 為 1 時，將跳過閘控評估和 MOE 以節省計算量。

參數：:

in_features (int) – 輸入的維度。
num_layers (int) – 模組中的層數。
low_rank (int) – 交叉矩陣的秩設定（預設值 = 1）。值必須始終 >= 1
activation (Union[torch.nn.Module, Callable[[torch.Tensor], torch.Tensor]]) – 定義專家時使用的非線性激活函數。預設值為 relu。

範例

batch_size = 3
num_layers = 2
in_features = 10
input = torch.randn(batch_size, in_features)
dcn = LowRankCrossNet(num_layers=num_layers, num_experts=5, low_rank=3)
output = dcn(input)

forward(input: Tensor) → Tensor¶

參數：:: input (torch.Tensor) – 形狀為 [batch_size, in_features] 的張量。
返回：:: 形狀為 [batch_size, in_features] 的張量。
返回類型：:: torch.Tensor

training: bool¶

class torchrec.modules.crossnet.VectorCrossNet(in_features: int, num_layers: int)¶

基於： Module

向量交叉網路 (Vector Cross Network) 可稱為 DCN-V1。

它也是一種特殊的低秩交叉網路，其中秩 = 1。在此版本中，在每一層上，我們不再保留兩個核心 W 和 V，而只保留一個向量核心 W (Nx1)。我們使用點積運算來計算特徵的「交叉」效應，從而節省了兩次矩陣乘法，進一步降低了計算成本並減少了可學習參數的數量。

在每一層 l 上，張量會被轉換為

\[x_{l+1} = x_0 * (W_l . x_l + b_l) + x_l\]

其中 \(W_l\) 是一個向量，\(*\) 表示元素乘法；\(.\) 表示點積運算。

參數：:

in_features (int) – 輸入的維度。
num_layers (int) – 模組中的層數。

範例

batch_size = 3
num_layers = 2
in_features = 10
input = torch.randn(batch_size, in_features)
dcn = VectorCrossNet(num_layers=num_layers)
output = dcn(input)

forward(input: Tensor) → Tensor¶

參數：:: input (torch.Tensor) – 形狀為 [batch_size, in_features] 的張量。
返回：:: 形狀為 [batch_size, in_features] 的張量。
返回類型：:: torch.Tensor

training: bool¶

torchrec.modules.deepfm¶

深度分解機模組

以下模組基於深度分解機 (DeepFM) 論文

DeepFM 類別實作了 DeepFM 架構
FactorizationMachine 類別實作了上述論文中提到的 FM。

class torchrec.modules.deepfm.DeepFM(dense_module: Module)¶

基於： Module

這是 DeepFM 模組

此模組不包含已發表論文的端到端功能。它僅涵蓋出版物的深度組件。它用於學習高階特徵交互作用。如果要學習低階特徵交互作用，請改用 FactorizationMachine 模組，它將共享此模組的相同嵌入輸入。

為了支援建模的靈活性，我們自訂了關鍵組件，例如

與公開論文不同的是，我們將輸入從原始稀疏特徵更改為特徵的嵌入。它允許嵌入維度和嵌入數量的靈活性，只要所有嵌入張量都具有相同的批次大小即可。
在公開論文的基礎上，我們允許使用者將隱藏層自訂為任何模組，而不僅限於 MLP。

該模組的總體架構如下

        1 x 10                  output
         /|\
          |                     pass into `dense_module`
          |
        1 x 90
         /|\
          |                     concat
          |
1 x 20, 1 x 30, 1 x 40          list of embeddings

參數：:: dense_module (nn.Module) – 任何可以在 DeepFM 中使用的自訂模組（例如 MLP）。此模組的 in_features 必須等於元素計數。例如，如果輸入嵌入為 [randn(3, 2, 3), randn(3, 4, 5)]，則 in_features 應為：2*3+4*5。

範例

import torch
from torchrec.fb.modules.deepfm import DeepFM
from torchrec.fb.modules.mlp import LazyMLP
batch_size = 3
output_dim = 30
# the input embedding are a torch.Tensor of [batch_size, num_embeddings, embedding_dim]
input_embeddings = [
    torch.randn(batch_size, 2, 64),
    torch.randn(batch_size, 2, 32),
]
dense_module = nn.Linear(192, output_dim)
deepfm = DeepFM(dense_module=dense_module)
deep_fm_output = deepfm(embeddings=input_embeddings)

forward(embeddings: List[Tensor]) → Tensor¶

參數：:

embeddings (List[torch.Tensor]) –

所有嵌入的列表（例如，密集、通用稀疏、專用稀疏、嵌入特徵、原始嵌入特徵），其形狀為

(batch_size, num_embeddings, embedding_dim)

為了便於操作，具有相同嵌入維度的嵌入可以選擇堆疊成單個張量。例如，當我們有 1 個訓練好的嵌入，維度為 32，5 個原生嵌入，維度為 64，以及 3 個密集特徵，維度為 16，我們可以準備嵌入列表，使其成為以下列表：

tensor(B, 1, 32) (trained_embedding with num_embeddings=1, embedding_dim=32)
tensor(B, 5, 64) (native_embedding with num_embeddings=5, embedding_dim=64)
tensor(B, 3, 16) (dense_features with num_embeddings=3, embedding_dim=32)

備註

所有輸入張量的 batch_size 必須相同。

返回：:

以扁平化和串聯的 embeddings 作為輸入的 dense_module 的輸出。

返回類型：:

torch.Tensor

training: bool¶

class torchrec.modules.deepfm.FactorizationMachine¶

基於： Module

這是 DeepFM 論文中提到的分解機模組，見 DeepFM 論文

此模組不涵蓋已發表論文的端到端功能。相反，它僅涵蓋出版物的 FM 部分，並用於學習二階特徵交互作用。

為了支援建模的靈活性，我們自訂了關鍵元件，使其與公開論文有所不同

我們將輸入從原始稀疏特徵更改為特徵的嵌入。這允許嵌入維度和嵌入數量的靈活性，只要所有嵌入張量具有相同的批次大小。

該模組的總體架構如下

        1 x 10                  output
         /|\
          |                     pass into `dense_module`
          |
        1 x 90
         /|\
          |                     concat
          |
1 x 20, 1 x 30, 1 x 40          list of embeddings

範例

batch_size = 3
# the input embedding are in torch.Tensor of [batch_size, num_embeddings, embedding_dim]
input_embeddings = [
    torch.randn(batch_size, 2, 64),
    torch.randn(batch_size, 2, 32),
]
fm = FactorizationMachine()
output = fm(embeddings=input_embeddings)

forward(embeddings: List[Tensor]) → Tensor¶

參數：:

embeddings (List[torch.Tensor]) –

所有嵌入的列表（例如，密集、通用稀疏、專用稀疏、嵌入特徵、原始嵌入特徵），其形狀為

(batch_size, num_embeddings, embedding_dim)

為了便於操作，具有相同嵌入維度的嵌入可以選擇堆疊成單個張量。例如，當我們有 1 個訓練好的嵌入，維度為 32，5 個原生嵌入，維度為 64，以及 3 個密集特徵，維度為 16，我們可以準備嵌入列表，使其成為以下列表：

tensor(B, 1, 32) (trained_embedding with num_embeddings=1, embedding_dim=32)
tensor(B, 5, 64) (native_embedding with num_embeddings=5, embedding_dim=64)
tensor(B, 3, 16) (dense_features with num_embeddings=3, embedding_dim=32)

備註

所有輸入張量的 batch_size 必須相同。

返回：:

以扁平化和串聯的 embeddings 作為輸入的 fm 輸出。預期為 [B, 1]。

返回類型：:

torch.Tensor

training: bool¶

torchrec.modules.embedding_configs¶

class torchrec.modules.embedding_configs.BaseEmbeddingConfig(num_embeddings: int, embedding_dim: int, name: str = '', data_type: torchrec.types.DataType = <DataType.FP32: 'FP32'>, feature_names: List[str] = <factory>, weight_init_max: Union[float, NoneType] = None, weight_init_min: Union[float, NoneType] = None, pruning_indices_remapping: Union[torch.Tensor, NoneType] = None, init_fn: Union[Callable[[torch.Tensor], Union[torch.Tensor, NoneType]], NoneType] = None, need_pos: bool = False)¶

基底： object

data_type: DataType = 'FP32'¶

embedding_dim: int¶

feature_names: List[str]¶

get_weight_init_max() → float¶

get_weight_init_min() → float¶

init_fn: Optional[Callable[[Tensor], Optional[Tensor]]] = None¶

name: str = ''¶

need_pos: bool = False¶

num_embeddings: int¶

num_features() → int¶

pruning_indices_remapping: Optional[Tensor] = None¶

weight_init_max: Optional[float] = None¶

weight_init_min: Optional[float] = None¶

class torchrec.modules.embedding_configs.EmbeddingBagConfig(num_embeddings: int, embedding_dim: int, name: str = '', data_type: torchrec.types.DataType = <DataType.FP32: 'FP32'>, feature_names: List[str] = <factory>, weight_init_max: Union[float, NoneType] = None, weight_init_min: Union[float, NoneType] = None, pruning_indices_remapping: Union[torch.Tensor, NoneType] = None, init_fn: Union[Callable[[torch.Tensor], Union[torch.Tensor, NoneType]], NoneType] = None, need_pos: bool = False, pooling: torchrec.modules.embedding_configs.PoolingType = <PoolingType.SUM: 'SUM'>)¶

基底： BaseEmbeddingConfig

pooling: PoolingType = 'SUM'¶

class torchrec.modules.embedding_configs.EmbeddingConfig(num_embeddings: int, embedding_dim: int, name: str = '', data_type: torchrec.types.DataType = <DataType.FP32: 'FP32'>, feature_names: List[str] = <factory>, weight_init_max: Union[float, NoneType] = None, weight_init_min: Union[float, NoneType] = None, pruning_indices_remapping: Union[torch.Tensor, NoneType] = None, init_fn: Union[Callable[[torch.Tensor], Union[torch.Tensor, NoneType]], NoneType] = None, need_pos: bool = False)¶

基底： BaseEmbeddingConfig

embedding_dim: int¶

feature_names: List[str]¶

num_embeddings: int¶

class torchrec.modules.embedding_configs.EmbeddingTableConfig(num_embeddings: int, embedding_dim: int, name: str = '', data_type: torchrec.types.DataType = <DataType.FP32: 'FP32'>, feature_names: List[str] = <factory>, weight_init_max: Union[float, NoneType] = None, weight_init_min: Union[float, NoneType] = None, pruning_indices_remapping: Union[torch.Tensor, NoneType] = None, init_fn: Union[Callable[[torch.Tensor], Union[torch.Tensor, NoneType]], NoneType] = None, need_pos: bool = False, pooling: torchrec.modules.embedding_configs.PoolingType = <PoolingType.SUM: 'SUM'>, is_weighted: bool = False, has_feature_processor: bool = False, embedding_names: List[str] = <factory>)¶

基底： BaseEmbeddingConfig

embedding_names: List[str]¶

has_feature_processor: bool = False¶

is_weighted: bool = False¶

pooling: PoolingType = 'SUM'¶

class torchrec.modules.embedding_configs.PoolingType(value)¶

基底： Enum

一個列舉。

MEAN = 'MEAN'¶

NONE = 'NONE'¶

SUM = 'SUM'¶

class torchrec.modules.embedding_configs.QuantConfig(activation, weight, per_table_weight_dtype)¶

基底： tuple

activation: PlaceholderObserver¶: 欄位編號 0 的別名

per_table_weight_dtype: Optional[Dict[str, dtype]]¶: 欄位編號 2 的別名

weight: PlaceholderObserver¶: 欄位編號 1 的別名

class torchrec.modules.embedding_configs.ShardingType(value)¶

基底： Enum

眾所周知的切片類型，由模組間優化使用。

COLUMN_WISE = 'column_wise'¶

DATA_PARALLEL = 'data_parallel'¶

ROW_WISE = 'row_wise'¶

TABLE_COLUMN_WISE = 'table_column_wise'¶

TABLE_ROW_WISE = 'table_row_wise'¶

TABLE_WISE = 'table_wise'¶

torchrec.modules.embedding_configs.data_type_to_dtype(data_type: DataType) → dtype¶

torchrec.modules.embedding_configs.data_type_to_sparse_type(data_type: DataType) → SparseType¶

torchrec.modules.embedding_configs.dtype_to_data_type(dtype: dtype) → DataType¶

torchrec.modules.embedding_configs.pooling_type_to_pooling_mode(pooling_type: PoolingType, sharding_type: Optional[ShardingType] = None) → PoolingMode¶

torchrec.modules.embedding_configs.pooling_type_to_str(pooling_type: PoolingType) → str¶

torchrec.modules.embedding_modules¶

class torchrec.modules.embedding_modules.EmbeddingBagCollection(tables: List[EmbeddingBagConfig], is_weighted: bool = False, device: Optional[device] = None)¶

Bases: EmbeddingBagCollectionInterface

EmbeddingBagCollection 代表一個池化嵌入（EmbeddingBags）的集合。

備註

EmbeddingBagCollection 是一個未分片的模組，並沒有針對效能進行優化。對於效能敏感的場景，請考慮使用分片版本 ShardedEmbeddingBagCollection。

它處理以 KeyedJaggedTensor 形式表示的稀疏數據，其值的格式為 [F X B X L]，其中

F：特徵（鍵）
B：批次大小
L：稀疏特徵的長度（不規則）

並輸出一個 KeyedTensor，其值的格式為 [B * (F * D)]，其中

F：特徵（鍵）
D：每個特徵（鍵）的嵌入維度
B：批次大小

參數：:

tables (List[EmbeddingBagConfig]) – 嵌入表的清單。
is_weighted (bool) – 輸入 KeyedJaggedTensor 是否為加權的。
device (Optional[torch.device]) – 預設計算裝置。

範例

table_0 = EmbeddingBagConfig(
    name="t1", embedding_dim=3, num_embeddings=10, feature_names=["f1"]
)
table_1 = EmbeddingBagConfig(
    name="t2", embedding_dim=4, num_embeddings=10, feature_names=["f2"]
)

ebc = EmbeddingBagCollection(tables=[table_0, table_1])

#        0       1        2  <-- batch
# "f1"   [0,1] None    [2]
# "f2"   [3]    [4]    [5,6,7]
#  ^
# feature

features = KeyedJaggedTensor(
    keys=["f1", "f2"],
    values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
    offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
)

pooled_embeddings = ebc(features)
print(pooled_embeddings.values())
tensor([[-0.8899, -0.1342, -1.9060, -0.0905, -0.2814, -0.9369, -0.7783],
    [ 0.0000,  0.0000,  0.0000,  0.1598,  0.0695,  1.3265, -0.1011],
    [-0.4256, -1.1846, -2.1648, -1.0893,  0.3590, -1.9784, -0.7681]],
    grad_fn=<CatBackward0>)
print(pooled_embeddings.keys())
['f1', 'f2']
print(pooled_embeddings.offset_per_key())
tensor([0, 3, 7])

property device: device¶

embedding_bag_configs() → List[EmbeddingBagConfig]¶

forward(features: KeyedJaggedTensor) → KeyedTensor¶

參數：:: features (KeyedJaggedTensor) – 格式為 [F X B X L] 的 KJT。
返回：:: KeyedTensor

is_weighted() → bool¶

reset_parameters() → None¶

training: bool¶

class torchrec.modules.embedding_modules.EmbeddingBagCollectionInterface(*args, **kwargs)¶

Bases: ABC, Module

用於 EmbeddingBagCollection 的介面。

abstract embedding_bag_configs() → List[EmbeddingBagConfig]¶

abstract forward(features: KeyedJaggedTensor) → KeyedTensor¶

定義每次呼叫時執行的計算。

應該被所有子類別覆蓋。

備註

雖然前向傳遞的配方需要在此函數中定義，但之後應該呼叫 Module 實例而不是這個，因為前者負責運行註冊的鉤子，而後者則默默地忽略它們。

abstract is_weighted() → bool¶

training: bool¶

class torchrec.modules.embedding_modules.EmbeddingCollection(tables: List[EmbeddingConfig], device: Optional[device] = None, need_indices: bool = False)¶

基底： EmbeddingCollectionInterface

EmbeddingCollection 代表非池化嵌入的集合。

備註

EmbeddingCollection 是一個未分片的模組，並未針對效能進行優化。對於效能敏感的場景，請考慮使用分片版本 ShardedEmbeddingCollection。

它處理 [F X B X L] 形式的 KeyedJaggedTensor 中的稀疏數據，其中

F：特徵（鍵）
B：批次大小
L：稀疏特徵的長度（變數）

並輸出 Dict[特徵 (鍵)，JaggedTensor]。每個 JaggedTensor 包含 (B * L) X D 形式的值，其中

B：批次大小
L：稀疏特徵的長度（不規則）
D：每個特徵（鍵）的嵌入維度，長度為 L 形式

參數：:

tables (List[EmbeddingConfig]) – 嵌入表的清單。
device (Optional[torch.device]) – 預設計算裝置。
need_indices (bool) – 我們是否需要將索引傳遞到最終的查找字典。

範例

e1_config = EmbeddingConfig(
    name="t1", embedding_dim=3, num_embeddings=10, feature_names=["f1"]
)
e2_config = EmbeddingConfig(
    name="t2", embedding_dim=3, num_embeddings=10, feature_names=["f2"]
)

ec = EmbeddingCollection(tables=[e1_config, e2_config])

#     0       1        2  <-- batch
# 0   [0,1] None    [2]
# 1   [3]    [4]    [5,6,7]
# ^
# feature

features = KeyedJaggedTensor.from_offsets_sync(
    keys=["f1", "f2"],
    values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
    offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
)
feature_embeddings = ec(features)
print(feature_embeddings['f2'].values())
tensor([[-0.2050,  0.5478,  0.6054],
[ 0.7352,  0.3210, -3.0399],
[ 0.1279, -0.1756, -0.4130],
[ 0.7519, -0.4341, -0.0499],
[ 0.9329, -1.0697, -0.8095]], grad_fn=<EmbeddingBackward>)

property device: device¶

embedding_configs() → List[EmbeddingConfig]¶

embedding_dim() → int¶

embedding_names_by_table() → List[List[str]]¶

forward(features: KeyedJaggedTensor) → Dict[str, JaggedTensor]¶

參數：:: features (KeyedJaggedTensor) – 格式為 [F X B X L] 的 KJT。
返回：:: Dict[str, JaggedTensor]

need_indices() → bool¶

reset_parameters() → None¶

training: bool¶

class torchrec.modules.embedding_modules.EmbeddingCollectionInterface(*args, **kwargs)¶

Bases: ABC, Module

EmbeddingCollection 的介面。

abstract embedding_configs() → List[EmbeddingConfig]¶

abstract embedding_dim() → int¶

abstract embedding_names_by_table() → List[List[str]]¶

abstract forward(features: KeyedJaggedTensor) → Dict[str, JaggedTensor]¶

定義每次呼叫時執行的計算。

應該被所有子類別覆蓋。

備註

雖然前向傳遞的配方需要在此函數中定義，但之後應該呼叫 Module 實例而不是這個，因為前者負責運行註冊的鉤子，而後者則默默地忽略它們。

abstract need_indices() → bool¶

training: bool¶

torchrec.modules.embedding_modules.get_embedding_names_by_table(tables: Union[List[EmbeddingBagConfig], List[EmbeddingConfig]]) → List[List[str]]¶

torchrec.modules.embedding_modules.process_pooled_embeddings(pooled_embeddings: List[Tensor], inverse_indices: Tensor) → Tensor¶

torchrec.modules.embedding_modules.reorder_inverse_indices(inverse_indices: Optional[Tuple[List[str], Tensor]], feature_names: List[str]) → Tensor¶

torchrec.modules.feature_processor¶

class torchrec.modules.feature_processor.BaseFeatureProcessor(*args, **kwargs)¶

基於： Module

特徵處理器的抽象基類。

abstract forward(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

定義每次呼叫時執行的計算。

應該被所有子類別覆蓋。

備註

雖然前向傳遞的配方需要在此函數中定義，但之後應該呼叫 Module 實例而不是這個，因為前者負責運行註冊的鉤子，而後者則默默地忽略它們。

training: bool¶

class torchrec.modules.feature_processor.BaseGroupedFeatureProcessor(*args, **kwargs)¶

基於： Module

分組特徵處理器的抽象基類

abstract forward(features: KeyedJaggedTensor) → KeyedJaggedTensor¶

定義每次呼叫時執行的計算。

應該被所有子類別覆蓋。

備註

雖然前向傳遞的配方需要在此函數中定義，但之後應該呼叫 Module 實例而不是這個，因為前者負責運行註冊的鉤子，而後者則默默地忽略它們。

training: bool¶

class torchrec.modules.feature_processor.PositionWeightedModule(max_feature_lengths: Dict[str, int], device: Optional[device] = None)¶

基礎類別： BaseFeatureProcessor

將位置權重添加到 ID 清單特徵。

參數：:: max_feature_lengths (Dict[str, int]) – 特徵名稱到 max_length 的映射。 max_length，又稱截斷大小，指定每個樣本具有的最大 ID 數量。對於每個特徵，其位置權重參數大小為 max_length。

forward(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

參數：:: features (Dict[str, JaggedTensor]) – 鍵到 JaggedTensor 的字典，表示特徵。
返回：:: 與輸入特徵相同，但填充了 weights 欄位。
返回類型：:: Dict[str, JaggedTensor]

reset_parameters() → None¶

training: bool¶

class torchrec.modules.feature_processor.PositionWeightedProcessor(max_feature_lengths: Dict[str, int], device: Optional[device] = None)¶

基礎類別： BaseGroupedFeatureProcessor

PositionWeightedProcessor 代表一個處理器，用於將位置權重應用於 KeyedJaggedTensor。

它可以處理未分片和分片的輸入，並輸出相應的輸出

參數：:

max_feature_lengths (Dict[str, int]) – 特徵長度字典，鍵是 feature_name，值是長度。
device (Optional[torch.device]) – 預設計算裝置。

範例

keys=["Feature0", "Feature1", "Feature2"]
values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 3, 4, 5, 6, 7])
lengths=torch.tensor([2, 0, 1, 1, 1, 3, 2, 3, 0])
features = KeyedJaggedTensor.from_lengths_sync(keys=keys, values=values, lengths=lengths)
pw = FeatureProcessorCollection(
    feature_processor_modules={key: PositionWeightedFeatureProcessor(max_feature_length=100) for key in keys}
)
result = pw(features)
# result is
# KeyedJaggedTensor({
#     "Feature0": {
#         "values": [[0, 1], [], [2]],
#         "weights": [[1.0, 1.0], [], [1.0]]
#     },
#     "Feature1": {
#         "values": [[3], [4], [5, 6, 7]],
#         "weights": [[1.0], [1.0], [1.0, 1.0, 1.0]]
#     },
#     "Feature2": {
#         "values": [[3, 4], [5, 6, 7], []],
#         "weights": [[1.0, 1.0], [1.0, 1.0, 1.0], []]
#     }
# })

forward(features: KeyedJaggedTensor) → KeyedJaggedTensor¶

在未分片或非流水線模型中，輸入特徵同時包含 fp_feature 和 non_fp_features，輸出將會過濾掉 non_fp 特徵。在分片流水線模型中，輸入特徵只能包含無或所有 feature_processed 特徵，因為輸入特徵來自 ebc 的 input_dist()，它會過濾掉不在 ebc 中的鍵。輸入大小與輸出大小相同

參數：:: features (KeyedJaggedTensor) – 輸入特徵
返回：:: KeyedJaggedTensor

named_buffers(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[Tuple[str, Tensor]]¶

傳回一個迭代器，迭代模組緩衝區，產生緩衝區的名稱和緩衝區本身。

參數：:

prefix (str) – 要預先添加到所有緩衝區名稱的前綴。
recurse (bool, optional) – 如果為 True，則產生此模組和所有子模組的緩衝區。否則，只產生此模組直接成員的緩衝區。預設值為 True。
remove_duplicate (bool, optional) – 是否移除結果中的重複緩衝區。預設值為 True。

產生：

(str, torch.Tensor) – 包含名稱和緩衝區的 Tuple

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

state_dict(destination: Optional[Dict[str, Any]] = None, prefix: str = '', keep_vars: bool = False) → Dict[str, Any]¶

傳回一個字典，其中包含對模組整體狀態的引用。

包括參數和持久緩衝區（例如，運行平均值）。鍵是相應的參數和緩衝區名稱。設定為 None 的參數和緩衝區不包括在內。

備註

傳回的物件是淺拷貝。它包含對模組參數和緩衝區的引用。

警告

目前 state_dict() 也接受 destination、prefix 和 keep_vars 的位置參數。但是，此功能已被棄用，未來版本將強制使用關鍵字參數。

警告

請避免使用參數 destination，因為它不是為最終用戶設計的。

參數：:

destination (dict, optional) – 如果提供，模組的狀態將更新到字典中，並傳回相同的物件。否則，將建立並傳回一個 OrderedDict。預設值：None。
prefix (str, optional) – 添加到參數和緩衝區名稱的前綴，用於組成 state_dict 中的鍵。預設值：''。
keep_vars (bool, optional) – 預設情況下，state dict 中傳回的 Tensor 與 autograd 分離。如果設定為 True，則不會執行分離。預設值：False。

返回：:

一個字典，包含模組的整體狀態

返回類型：:

dict

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> module.state_dict().keys()
['bias', 'weight']

training: bool¶

torchrec.modules.feature_processor.offsets_to_range_traceble(offsets: Tensor, values: Tensor) → Tensor¶

torchrec.modules.feature_processor.position_weighted_module_update_features(features: Dict[str, JaggedTensor], weighted_features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

torchrec.modules.lazy_extension¶

class torchrec.modules.lazy_extension.LazyModuleExtensionMixin(*args, **kwargs)¶

基底：LazyModuleMixin

這是 LazyModuleMixin 的臨時擴充，用於支援將關鍵字引數傳遞給懶惰模組的 forward 方法。

長期計畫是將此功能上推到 LazyModuleMixin。如需詳細資訊，請參閱 https://github.com/pytorch/pytorch/issues/59923。

請參閱 TestLazyModuleExtensionMixin，其中包含確保以下條件的單元測試：

LazyModuleExtensionMixin._infer_parameters 與 torch.nn.modules.lazy.LazyModuleMixin._infer_parameters 具有原始碼同位性，但前者可以接受關鍵字引數。
LazyModuleExtensionMixin._call_impl 與 torch.nn.Module._call_impl 具有原始碼同位性，但前者可以將關鍵字引數傳遞給前向預先鉤子。

apply(fn: Callable[[Module], None]) → Module¶

遞迴地將 fn 應用於每個子模組（由 .children() 返回）以及自身。典型用途包括初始化模型的參數。

備註

在未初始化的懶惰模組上呼叫 apply() 將導致錯誤。在對懶惰模組呼叫 apply() 之前，使用者需要先初始化懶惰模組（透過執行虛擬前向傳遞）。

參數：:: fn (torch.nn.Module -> None) – 要應用於每個子模組的函數。
返回：:: 自身
返回類型：:: torch.nn.Module

範例

@torch.no_grad()
def init_weights(m):
    print(m)
    if type(m) == torch.nn.LazyLinear:
        m.weight.fill_(1.0)
        print(m.weight)

linear = torch.nn.LazyLinear(2)
linear.apply(init_weights)  # this fails, because `linear` (a lazy-module) hasn't been initialized yet

input = torch.randn(2, 10)
linear(input)  # run a dummy forward pass to initialize the lazy-module

linear.apply(init_weights)  # this works now

torchrec.modules.lazy_extension.lazy_apply(module: Module, fn: Callable[[Module], None]) → Module¶

將函數附加到模組，該函數將在第一次前向傳遞後（即在所有子模組和參數初始化之後）遞迴地應用於模組的每個子模組（由 .children() 返回）以及模組自身。

典型用途包括初始化懶惰模組（即從 LazyModuleMixin 繼承的模組）的參數數值。

備註

lazy_apply() 可以用於懶惰和非懶惰模組。

參數：:

module (torch.nn.Module) – 要遞迴應用 fn 的模組。
fn (Callable[[torch.nn.Module], None]) – 要附加到 module 並稍後應用於 module 的每個子模組和 module 自身的函數。

返回：:

附加了 fn 的 module。

返回類型：:

torch.nn.Module

範例

@torch.no_grad()
def init_weights(m):
    print(m)
    if type(m) == torch.nn.LazyLinear:
        m.weight.fill_(1.0)
        print(m.weight)

linear = torch.nn.LazyLinear(2)
lazy_apply(linear, init_weights)  # doesn't run `init_weights` immediately
input = torch.randn(2, 10)
linear(input)  # runs `init_weights` only once, right after first forward pass

seq = torch.nn.Sequential(torch.nn.LazyLinear(2), torch.nn.LazyLinear(2))
lazy_apply(seq, init_weights)  # doesn't run `init_weights` immediately
input = torch.randn(2, 10)
seq(input)  # runs `init_weights` only once, right after first forward pass

torchrec.modules.mlp¶

class torchrec.modules.mlp.MLP(in_size: int, layer_sizes: ~typing.List[int], bias: bool = True, activation: ~typing.Union[str, ~typing.Callable[[], ~torch.nn.modules.module.Module], ~torch.nn.modules.module.Module, ~typing.Callable[[~torch.Tensor], ~torch.Tensor]] = <built-in method relu of type object>, device: ~typing.Optional[~torch.device] = None, dtype: ~torch.dtype = torch.float32)¶

基於： Module

依序應用一組感知器模組（即多層感知器）。

參數：:

in_size (int) – 輸入的 in_size。
layer_sizes (List[int]) – 每個感知器模組的 out_size。
bias (bool) – 如果設定為 False，則該層不會學習加性偏差。預設值：True。
activation (str, Union[Callable[[], torch.nn.Module], torch.nn.Module, Callable[[torch.Tensor], torch.Tensor]]) – 要應用於每個感知器模組的線性變換輸出的激活函數。如果 activation 是 str，我們目前僅支援以下字串： “relu”、“sigmoid” 和 “swish_layernorm”。如果 activation 是 Callable[[], torch.nn.Module]，則每個感知器模組將呼叫一次 activation() 以產生該感知器模組的激活模組，並且這些激活模組之間不會共享參數。一個用例是當所有激活模組共享相同的建構函數引數，但不共享實際的模組參數時。預設值：torch.relu。
device (Optional[torch.device]) – 預設計算裝置。

範例

batch_size = 3
in_size = 40
input = torch.randn(batch_size, in_size)

layer_sizes = [16, 8, 4]
mlp_module = MLP(in_size, layer_sizes, bias=True)
output = mlp_module(input)
assert list(output.shape) == [batch_size, layer_sizes[-1]]

forward(input: Tensor) → Tensor¶

參數：:: input (torch.Tensor) – 形狀為 (B, I) 的張量，其中 I 是每個輸入樣本中的元素數量。
返回：:: 形狀為 (B, O) 的張量，其中 O 是最後一個感知器模組的 out_size。
返回類型：:: torch.Tensor

training: bool¶

class torchrec.modules.mlp.Perceptron(in_size: int, out_size: int, bias: bool = True, activation: ~typing.Union[~torch.nn.modules.module.Module, ~typing.Callable[[~torch.Tensor], ~torch.Tensor]] = <built-in method relu of type object>, device: ~typing.Optional[~torch.device] = None, dtype: ~torch.dtype = torch.float32)¶

基於： Module

應用線性變換和激活。

參數：:

in_size (int) – 每個輸入樣本中的元素數量。
out_size (int) – 每個輸出樣本中的元素數量。
bias (bool) – 如果設定為 False，則該層不會學習加性偏差。預設值：True。
activation (Union[torch.nn.Module, Callable[[torch.Tensor], torch.Tensor]]) – 要應用於線性變換輸出的激活函數。預設值：torch.relu。
device (Optional[torch.device]) – 預設計算裝置。

範例

batch_size = 3
in_size = 40
input = torch.randn(batch_size, in_size)

out_size = 16
perceptron = Perceptron(in_size, out_size, bias=True)

output = perceptron(input)
assert list(output) == [batch_size, out_size]

forward(input: Tensor) → Tensor¶

參數：:

input (torch.Tensor) – 形狀為 (B, I) 的張量，其中 I 是每個輸入樣本中的元素數量。

返回：:

形狀為 (B, O) 的張量，其中 O 是每個: 輸出樣本中每個通道的元素數量（即 out_size）。

返回類型：:

torch.Tensor

training: bool¶

torchrec.modules.utils¶

class torchrec.modules.utils.SequenceVBEContext(recat: torch.Tensor, unpadded_lengths: torch.Tensor, reindexed_lengths: torch.Tensor, reindexed_length_per_key: List[int], reindexed_values: Union[torch.Tensor, NoneType] = None)¶

基底：Multistreamable

recat: Tensor¶

record_stream(stream: Stream) → None¶: 請見 https://pytorch.dev.org.tw/docs/stable/generated/torch.Tensor.record_stream.html

reindexed_length_per_key: List[int]¶

reindexed_lengths: Tensor¶

reindexed_values: Optional[Tensor] = None¶

unpadded_lengths: Tensor¶

torchrec.modules.utils.check_module_output_dimension(module: Union[Iterable[Module], Module], in_features: int, out_features: int) → bool¶: 驗證給定模組或模組列表的 out_features 是否與指定的數量相符。如果給定的是模組列表或 ModuleList，則遞迴檢查所有子模組。

torchrec.modules.utils.construct_jagged_tensors(embeddings: Tensor, features: KeyedJaggedTensor, embedding_names: List[str], need_indices: bool = False, features_to_permute_indices: Optional[Dict[str, List[int]]] = None, original_features: Optional[KeyedJaggedTensor] = None, reverse_indices: Optional[Tensor] = None, seq_vbe_ctx: Optional[SequenceVBEContext] = None) → Dict[str, JaggedTensor]¶

torchrec.modules.utils.construct_jagged_tensors_inference(embeddings: Tensor, lengths: Tensor, values: Tensor, embedding_names: List[str], need_indices: bool = False, features_to_permute_indices: Optional[Dict[str, List[int]]] = None, reverse_indices: Optional[Tensor] = None) → Dict[str, JaggedTensor]¶

torchrec.modules.utils.construct_modulelist_from_single_module(module: Module, sizes: Tuple[int, ...]) → Module¶: 給定單一模組，透過複製提供的模組並重新初始化線性層，建構大小為 sizes 的（巢狀）ModuleList。

torchrec.modules.utils.convert_list_of_modules_to_modulelist(modules: Iterable[Module], sizes: Tuple[int, ...]) → Module¶

torchrec.modules.utils.deterministic_dedup(ids: Tensor) → Tuple[Tensor, Tensor]¶: 為了消除衝突更新中的競爭條件，請移除重複的 ID。只會保留最後一個存在的重複 ID。傳回排序後的唯一 ID 和最後一個存在的位置

torchrec.modules.utils.extract_module_or_tensor_callable(module_or_callable: Union[Callable[[], Module], Module, Callable[[Tensor], Tensor]]) → Union[Module, Callable[[Tensor], Tensor]]¶

torchrec.modules.utils.get_module_output_dimension(module: Union[Callable[[Tensor], Tensor], Module], in_features: int) → int¶

torchrec.modules.utils.init_mlp_weights_xavier_uniform(m: Module) → None¶

torchrec.modules.utils.jagged_index_select_with_empty(values: Tensor, ids: Tensor, offsets: Tensor, output_offsets: Tensor) → Tensor¶

torchrec.modules.mc_modules¶

class torchrec.modules.mc_modules.DistanceLFU_EvictionPolicy(decay_exponent: float = 1.0, threshold_filtering_func: Optional[Callable[[Tensor], Tuple[Tensor, Union[float, Tensor]]]] = None)¶

基底： MCHEvictionPolicy

coalesce_history_metadata(current_iter: int, history_metadata: Dict[str, Tensor], unique_ids_counts: Tensor, unique_inverse_mapping: Tensor, additional_ids: Optional[Tensor] = None, threshold_mask: Optional[Tensor] = None) → Dict[str, Tensor]¶

參數：history_metadata (Dict[str, torch.Tensor])：歷史中繼資料字典 additional_ids (torch.Tensor)：要用作歷史記錄一部分的額外 ID unique_inverse_mapping (torch.Tensor)：從

torch.cat[history_accumulator, additional_ids] 生成的 torch.unique 反向映射。用於將歷史中繼資料張量索引映射到它們的合併張量索引。

合併中繼資料歷史記錄緩衝區，並傳回已處理中繼資料張量的字典。

property metadata_info: List[MCHEvictionPolicyMetadataInfo]¶

record_history_metadata(current_iter: int, incoming_ids: Tensor, history_metadata: Dict[str, Tensor]) → None¶

參數： current_iter (int): 當前迭代 incoming_ids (torch.Tensor): 進入的 ID history_metadata (Dict[str, torch.Tensor]): 歷史中繼資料字典

根據進入的 ID 計算並記錄中繼資料: 以供已實作的逐出策略使用。

update_metadata_and_generate_eviction_scores(current_iter: int, mch_size: int, coalesced_history_argsort_mapping: Tensor, coalesced_history_sorted_unique_ids_counts: Tensor, coalesced_history_mch_matching_elements_mask: Tensor, coalesced_history_mch_matching_indices: Tensor, mch_metadata: Dict[str, Tensor], coalesced_history_metadata: Dict[str, Tensor]) → Tuple[Tensor, Tensor]¶

參數

傳回 (evicted_indices, selected_new_indices) 的 Tuple，其中: evicted_indices 是要逐出的 mch map 中的索引，而 selected_new_indices 是要加入 mch 的合併歷史記錄中 ID 的索引。

class torchrec.modules.mc_modules.LFU_EvictionPolicy(threshold_filtering_func: Optional[Callable[[Tensor], Tuple[Tensor, Union[float, Tensor]]]] = None)¶

基底： MCHEvictionPolicy

coalesce_history_metadata(current_iter: int, history_metadata: Dict[str, Tensor], unique_ids_counts: Tensor, unique_inverse_mapping: Tensor, additional_ids: Optional[Tensor] = None, threshold_mask: Optional[Tensor] = None) → Dict[str, Tensor]¶

參數：history_metadata (Dict[str, torch.Tensor])：歷史中繼資料字典 additional_ids (torch.Tensor)：要用作歷史記錄一部分的額外 ID unique_inverse_mapping (torch.Tensor)：從

torch.cat[history_accumulator, additional_ids] 生成的 torch.unique 反向映射。用於將歷史中繼資料張量索引映射到它們的合併張量索引。

合併中繼資料歷史記錄緩衝區，並傳回已處理中繼資料張量的字典。

property metadata_info: List[MCHEvictionPolicyMetadataInfo]¶

record_history_metadata(current_iter: int, incoming_ids: Tensor, history_metadata: Dict[str, Tensor]) → None¶

參數： current_iter (int): 當前迭代 incoming_ids (torch.Tensor): 進入的 ID history_metadata (Dict[str, torch.Tensor]): 歷史中繼資料字典

根據進入的 ID 計算並記錄中繼資料: 以供已實作的逐出策略使用。

update_metadata_and_generate_eviction_scores(current_iter: int, mch_size: int, coalesced_history_argsort_mapping: Tensor, coalesced_history_sorted_unique_ids_counts: Tensor, coalesced_history_mch_matching_elements_mask: Tensor, coalesced_history_mch_matching_indices: Tensor, mch_metadata: Dict[str, Tensor], coalesced_history_metadata: Dict[str, Tensor]) → Tuple[Tensor, Tensor]¶

參數

傳回 (evicted_indices, selected_new_indices) 的 Tuple，其中: evicted_indices 是要逐出的 mch map 中的索引，而 selected_new_indices 是要加入 mch 的合併歷史記錄中 ID 的索引。

class torchrec.modules.mc_modules.LRU_EvictionPolicy(decay_exponent: float = 1.0, threshold_filtering_func: Optional[Callable[[Tensor], Tuple[Tensor, Union[float, Tensor]]]] = None)¶

基底： MCHEvictionPolicy

coalesce_history_metadata(current_iter: int, history_metadata: Dict[str, Tensor], unique_ids_counts: Tensor, unique_inverse_mapping: Tensor, additional_ids: Optional[Tensor] = None, threshold_mask: Optional[Tensor] = None) → Dict[str, Tensor]¶

參數：history_metadata (Dict[str, torch.Tensor])：歷史中繼資料字典 additional_ids (torch.Tensor)：要用作歷史記錄一部分的額外 ID unique_inverse_mapping (torch.Tensor)：從

torch.cat[history_accumulator, additional_ids] 生成的 torch.unique 反向映射。用於將歷史中繼資料張量索引映射到它們的合併張量索引。

合併中繼資料歷史記錄緩衝區，並傳回已處理中繼資料張量的字典。

property metadata_info: List[MCHEvictionPolicyMetadataInfo]¶

record_history_metadata(current_iter: int, incoming_ids: Tensor, history_metadata: Dict[str, Tensor]) → None¶

參數： current_iter (int): 當前迭代 incoming_ids (torch.Tensor): 進入的 ID history_metadata (Dict[str, torch.Tensor]): 歷史中繼資料字典

根據進入的 ID 計算並記錄中繼資料: 以供已實作的逐出策略使用。

update_metadata_and_generate_eviction_scores(current_iter: int, mch_size: int, coalesced_history_argsort_mapping: Tensor, coalesced_history_sorted_unique_ids_counts: Tensor, coalesced_history_mch_matching_elements_mask: Tensor, coalesced_history_mch_matching_indices: Tensor, mch_metadata: Dict[str, Tensor], coalesced_history_metadata: Dict[str, Tensor]) → Tuple[Tensor, Tensor]¶

參數

傳回 (evicted_indices, selected_new_indices) 的 Tuple，其中: evicted_indices 是要逐出的 mch map 中的索引，而 selected_new_indices 是要加入 mch 的合併歷史記錄中 ID 的索引。

class torchrec.modules.mc_modules.MCHEvictionPolicy(metadata_info: List[MCHEvictionPolicyMetadataInfo], threshold_filtering_func: Optional[Callable[[Tensor], Tuple[Tensor, Union[float, Tensor]]]] = None)¶

基底： ABC

abstract coalesce_history_metadata(current_iter: int, history_metadata: Dict[str, Tensor], unique_ids_counts: Tensor, unique_inverse_mapping: Tensor, additional_ids: Optional[Tensor] = None, threshold_mask: Optional[Tensor] = None) → Dict[str, Tensor]¶

參數：history_metadata (Dict[str, torch.Tensor])：歷史中繼資料字典 additional_ids (torch.Tensor)：要用作歷史記錄一部分的額外 ID unique_inverse_mapping (torch.Tensor)：從

torch.cat[history_accumulator, additional_ids] 生成的 torch.unique 反向映射。用於將歷史中繼資料張量索引映射到它們的合併張量索引。

合併中繼資料歷史記錄緩衝區，並傳回已處理中繼資料張量的字典。

abstract property metadata_info: List[MCHEvictionPolicyMetadataInfo]¶

abstract record_history_metadata(current_iter: int, incoming_ids: Tensor, history_metadata: Dict[str, Tensor]) → None¶

參數： current_iter (int): 當前迭代 incoming_ids (torch.Tensor): 進入的 ID history_metadata (Dict[str, torch.Tensor]): 歷史中繼資料字典

根據進入的 ID 計算並記錄中繼資料: 以供已實作的逐出策略使用。

abstract update_metadata_and_generate_eviction_scores(current_iter: int, mch_size: int, coalesced_history_argsort_mapping: Tensor, coalesced_history_sorted_unique_ids_counts: Tensor, coalesced_history_mch_matching_elements_mask: Tensor, coalesced_history_mch_matching_indices: Tensor, mch_metadata: Dict[str, Tensor], coalesced_history_metadata: Dict[str, Tensor]) → Tuple[Tensor, Tensor]¶

參數

傳回 (evicted_indices, selected_new_indices) 的 Tuple，其中: evicted_indices 是要逐出的 mch map 中的索引，而 selected_new_indices 是要加入 mch 的合併歷史記錄中 ID 的索引。

class torchrec.modules.mc_modules.MCHEvictionPolicyMetadataInfo(metadata_name, is_mch_metadata, is_history_metadata)¶

基底： tuple

is_history_metadata: bool¶: 欄位編號 2 的別名

is_mch_metadata: bool¶: 欄位編號 1 的別名

metadata_name: str¶: 欄位編號 0 的別名

class torchrec.modules.mc_modules.MCHManagedCollisionModule(zch_size: int, device: device, eviction_policy: MCHEvictionPolicy, eviction_interval: int, input_hash_size: int = 9223372036854775808, input_hash_func: Optional[Callable[[Tensor, int], Tensor]] = None, mch_size: Optional[int] = None, mch_hash_func: Optional[Callable[[Tensor, int], Tensor]] = None, name: Optional[str] = None, output_global_offset: int = 0)¶

基底： ManagedCollisionModule

ZCH / MCH 受管衝突模組

參數：:

zch_size (int) – 輸出 ID 範圍，在 [output_size_offset, output_size_offset + zch_size - 1) 之內
device (torch.device) – 將執行此模組的裝置
eviction_policy (eviction 政策) – 要使用的 eviction 政策
eviction_interval (int) – 觸發 eviction 政策的間隔
input_hash_size (int) – 輸入特徵 ID 範圍，將作為第二個參數傳遞給 input_hash_func
input_hash_func (Optional[Callable]) – 用於為輸入特徵產生雜湊的函數。此函數通常用於驅動相同或大於輸入數據範圍的均勻分佈
mch_size (Optional[int]) – 剩餘輸出的大小（即傳統 MCH），實驗性功能。ID 在內部由 output_size_offset + zch_output_range 移動
mch_hash_func (Optional[Callable]) – 用於為剩餘特徵產生雜湊的函數。將雜湊到 mch_size。
output_global_offset (int) – 輸出範圍的輸出 ID 位移量，通常僅在分片應用中使用。

evict() → Optional[Tensor]¶: 如果在此迭代中不應進行逐出，則返回 None。否則，返回要重置的插槽 ID。逐出時，此模組應針對這些插槽重置其狀態，並假設下游模組將妥善處理。

forward(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶: 參數：feature (JaggedTensor]): 特徵表示 :returns: 修改後的 JT :rtype: Dict[str, JaggedTensor]

input_size() → int¶: 傳回輸入的數值範圍，用於分片資訊

output_size() → int¶: 傳回輸出的數值範圍，用於驗證與下游嵌入查找

preprocess(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

profile(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

rebuild_with_output_id_range(output_id_range: Tuple[int, int], device: Optional[device] = None) → MCHManagedCollisionModule¶: 用於建立用於 RW 分片的本機 MC 模組，目前是權宜之計

remap(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

training: bool¶

class torchrec.modules.mc_modules.ManagedCollisionCollection(managed_collision_modules: Dict[str, ManagedCollisionModule], embedding_configs: List[BaseEmbeddingConfig])¶

基於： Module

ManagedCollisionCollection 表示受控衝突模組的集合。傳遞到 MCC 的輸入將由受控衝突模組重新映射

並返回。

參數：:

managed_collision_modules (Dict[str, ManagedCollisionModule]) – 受控衝突模組的字典
embedding_confgs (List[BaseEmbeddingConfig]) – 嵌入配置的清單，適用於具有受控衝突模組的每個表格

embedding_configs() → List[BaseEmbeddingConfig]¶

evict() → Dict[str, Optional[Tensor]]¶

forward(features: KeyedJaggedTensor) → KeyedJaggedTensor¶

定義每次呼叫時執行的計算。

應該被所有子類別覆蓋。

備註

雖然前向傳遞的配方需要在此函數中定義，但之後應該呼叫 Module 實例而不是這個，因為前者負責運行註冊的鉤子，而後者則默默地忽略它們。

training: bool¶

class torchrec.modules.mc_modules.ManagedCollisionModule(device: device)¶

基於： Module

ManagedCollisionModule 的抽象基類。將輸入 ID 映射到範圍 [0, max_output_id)。

參數：:

max_output_id (int) – 重新映射 ID 的最大輸出值。
input_hash_size (int) – 輸入範圍的最大值，即 [0, input_hash_size)
remapping_range_start_index (int) – 重新映射範圍的相對起始索引
device (torch.device) – 預設運算裝置。

範例：: jt = JaggedTensor(…) mcm = ManagedCollisionModule(…) mcm_jt = mcm(fp)

property device: device¶

abstract evict() → Optional[Tensor]¶: 如果在此迭代中不應進行逐出，則返回 None。否則，返回要重置的插槽 ID。逐出時，此模組應針對這些插槽重置其狀態，並假設下游模組將妥善處理。

abstract forward(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

定義每次呼叫時執行的計算。

應該被所有子類別覆蓋。

備註

雖然前向傳遞的配方需要在此函數中定義，但之後應該呼叫 Module 實例而不是這個，因為前者負責運行註冊的鉤子，而後者則默默地忽略它們。

abstract input_size() → int¶: 傳回輸入的數值範圍，用於分片資訊

abstract output_size() → int¶: 傳回輸出的數值範圍，用於驗證與下游嵌入查找

abstract preprocess(features: Dict[str, JaggedTensor]) → Dict[str, JaggedTensor]¶

abstract rebuild_with_output_id_range(output_id_range: Tuple[int, int], device: Optional[device] = None) → ManagedCollisionModule¶: 用於建立用於 RW 分片的本機 MC 模組，目前是權宜之計

training: bool¶

torchrec.modules.mc_modules.apply_mc_method_to_jt_dict(method: str, features_dict: Dict[str, JaggedTensor], table_to_features: Dict[str, List[str]], managed_collisions: ModuleDict) → Dict[str, JaggedTensor]¶: 將 MC 方法應用於 JaggedTensors 的字典，返回具有相同順序的更新後字典

torchrec.modules.mc_modules.average_threshold_filter(id_counts: Tensor) → Tuple[Tensor, Tensor]¶: 閾值為 id_counts 的平均值。如果 id 的計數嚴格大於平均值，則會添加該 id。

torchrec.modules.mc_modules.dynamic_threshold_filter(id_counts: Tensor, threshold_skew_multiplier: float = 10.0) → Tuple[Tensor, Tensor]¶: 閾值為 total_count / num_ids * threshold_skew_multiplier。如果 id 的計數嚴格大於閾值，則會添加該 id。

torchrec.modules.mc_modules.probabilistic_threshold_filter(id_counts: Tensor, per_id_probability: float = 0.01) → Tuple[Tensor, Tensor]¶: 每個 id 都有 per_id_probability 的機率被添加。例如，如果 per_id_probability 為 0.01 且 id 出現 100 次，則它有 60% 的機率被添加。更精確地說，id 分數為 1 - (1 - per_id_probability) ^ id_count，對於隨機生成的閾值，id 分數是它被添加的機率。

torchrec.modules.mc_embedding_modules¶

class torchrec.modules.mc_embedding_modules.BaseManagedCollisionEmbeddingCollection(embedding_module: Union[EmbeddingBagCollection, EmbeddingCollection], managed_collision_collection: ManagedCollisionCollection, return_remapped_features: bool = False)¶

基於： Module

BaseManagedCollisionEmbeddingCollection 表示一個 EC/EBC 模組和一組受管衝突模組。在傳遞到嵌入集合之前，MC-EC/EBC 的輸入將首先由受管衝突模組修改。

參數：:

embedding_module – 用於查詢嵌入的 EmbeddingCollection
managed_collision_modules – 受管衝突模組的字典
return_remapped_features (bool) – 是否除了嵌入之外還返回重新映射的輸入特徵

forward(features: KeyedJaggedTensor) → Tuple[Union[KeyedTensor, Dict[str, JaggedTensor]], Optional[KeyedJaggedTensor]]¶

定義每次呼叫時執行的計算。

應該被所有子類別覆蓋。

備註

雖然前向傳遞的配方需要在此函數中定義，但之後應該呼叫 Module 實例而不是這個，因為前者負責運行註冊的鉤子，而後者則默默地忽略它們。

training: bool¶

class torchrec.modules.mc_embedding_modules.ManagedCollisionEmbeddingBagCollection(embedding_bag_collection: EmbeddingBagCollection, managed_collision_collection: ManagedCollisionCollection, return_remapped_features: bool = False)¶

基底： BaseManagedCollisionEmbeddingCollection

ManagedCollisionEmbeddingBagCollection 表示一個 EmbeddingBagCollection 模組和一組受管衝突模組。在傳遞到嵌入袋集合之前，MC-EBC 的輸入將首先由受管衝突模組修改。

有關輸入和輸出類型的詳細信息，請參閱 EmbeddingBagCollection

參數：:

embedding_module – 用於查詢嵌入的 EmbeddingBagCollection
managed_collision_modules – 受管衝突模組的字典
return_remapped_features (bool) – 是否除了嵌入之外還返回重新映射的輸入特徵

training: bool¶

class torchrec.modules.mc_embedding_modules.ManagedCollisionEmbeddingCollection(embedding_collection: EmbeddingCollection, managed_collision_collection: ManagedCollisionCollection, return_remapped_features: bool = False)¶

基底： BaseManagedCollisionEmbeddingCollection

ManagedCollisionEmbeddingCollection 表示一個 EmbeddingCollection 模組和一組受管衝突模組。在傳遞到嵌入集合之前，MC-EC 的輸入將首先由受管衝突模組修改。

有關輸入和輸出類型的詳細信息，請參閱 EmbeddingCollection

參數：:

embedding_module – 用於查詢嵌入的 EmbeddingCollection
managed_collision_modules – 受管衝突模組的字典
return_remapped_features (bool) – 是否除了嵌入之外還返回重新映射的輸入特徵

training: bool¶

torchrec.modules.mc_embedding_modules.evict(evictions: Dict[str, Optional[Tensor]], ebc: Union[EmbeddingBagCollection, EmbeddingCollection]) → None¶

torchrec.modules¶

torchrec.modules.activation¶

torchrec.modules.crossnet¶

torchrec.modules.deepfm¶

torchrec.modules.embedding_configs¶

torchrec.modules.embedding_modules¶

torchrec.modules.feature_processor¶

torchrec.modules.lazy_extension¶

torchrec.modules.mlp¶

torchrec.modules.utils¶

torchrec.modules.mc_modules¶

torchrec.modules.mc_embedding_modules¶

文件

教學課程

資源