模块¶

标准的 TorchRec 模块代表嵌入表的集合

EmbeddingBagCollection 是 torch.nn.EmbeddingBag 的集合
EmbeddingCollection 是 torch.nn.Embedding 的集合

这些模块是通过标准化的配置类构建的

EmbeddingBagConfig 用于 EmbeddingBagCollection
EmbeddingConfig 用于 EmbeddingCollection

class torchrec.modules.embedding_configs.EmbeddingBagConfig(num_embeddings: int, embedding_dim: int, name: str = '', data_type: ~torchrec.types.DataType = DataType.FP32, feature_names: ~typing.List[str] = <factory>, weight_init_max: ~typing.Optional[float] = None, weight_init_min: ~typing.Optional[float] = None, num_embeddings_post_pruning: ~typing.Optional[int] = None, init_fn: ~typing.Optional[~typing.Callable[[~torch.Tensor], ~typing.Optional[~torch.Tensor]]] = None, need_pos: bool = False, input_dim: ~typing.Optional[int] = None, total_num_buckets: ~typing.Optional[int] = None, use_virtual_table: bool = False, virtual_table_eviction_policy: ~typing.Optional[~torchrec.modules.embedding_configs.VirtualTableEvictionPolicy] = None, pooling: ~torchrec.modules.embedding_configs.PoolingType = PoolingType.SUM)¶

Bases: BaseEmbeddingConfig

EmbeddingBagConfig 是一个数据类，用于表示单个嵌入表，其输出旨在被池化。

参数:: pooling (PoolingType) – 池化类型。

class torchrec.modules.embedding_configs.EmbeddingConfig(num_embeddings: int, embedding_dim: int, name: str = '', data_type: ~torchrec.types.DataType = DataType.FP32, feature_names: ~typing.List[str] = <factory>, weight_init_max: ~typing.Optional[float] = None, weight_init_min: ~typing.Optional[float] = None, num_embeddings_post_pruning: ~typing.Optional[int] = None, init_fn: ~typing.Optional[~typing.Callable[[~torch.Tensor], ~typing.Optional[~torch.Tensor]]] = None, need_pos: bool = False, input_dim: ~typing.Optional[int] = None, total_num_buckets: ~typing.Optional[int] = None, use_virtual_table: bool = False, virtual_table_eviction_policy: ~typing.Optional[~torchrec.modules.embedding_configs.VirtualTableEvictionPolicy] = None)¶

Bases: BaseEmbeddingConfig

EmbeddingConfig 是一个数据类，用于表示单个嵌入表。

class torchrec.modules.embedding_configs.BaseEmbeddingConfig(num_embeddings: int, embedding_dim: int, name: str = '', data_type: ~torchrec.types.DataType = DataType.FP32, feature_names: ~typing.List[str] = <factory>, weight_init_max: ~typing.Optional[float] = None, weight_init_min: ~typing.Optional[float] = None, num_embeddings_post_pruning: ~typing.Optional[int] = None, init_fn: ~typing.Optional[~typing.Callable[[~torch.Tensor], ~typing.Optional[~torch.Tensor]]] = None, need_pos: bool = False, input_dim: ~typing.Optional[int] = None, total_num_buckets: ~typing.Optional[int] = None, use_virtual_table: bool = False, virtual_table_eviction_policy: ~typing.Optional[~torchrec.modules.embedding_configs.VirtualTableEvictionPolicy] = None)¶

嵌入配置的基类。

参数:

num_embeddings (int) – 嵌入数量。
embedding_dim (int) – 嵌入维度。
name (str) – 嵌入表的名称。
data_type (DataType) – 嵌入表的数据类型。
feature_names (List[str]) – 特征名称列表。
weight_init_max (Optional[float]) – 权重初始化的最大值。
weight_init_min (Optional[float]) – 权重初始化的最小值。
num_embeddings_post_pruning (Optional[int]) – 剪枝后的嵌入数量（用于推理）。如果为 None，则不进行剪枝。
init_fn (Optional[Callable[[torch.Tensor], Optional[torch.Tensor]]]) – 嵌入权重的初始化函数。
need_pos (bool) – 表是否按位置加权。
total_num_buckets (Optional[int]) – 全局桶的数量，在模型生命周期内保持不变
use_virtual_table (bool) – 表是否使用虚拟空间（例如 2^50 的数量级）来存储嵌入内存的指示符，虚拟表的内存是动态的，仅在 id 被训练时物化。这需要与 EmbeddingComputeKernel 中的 SSD/DRAM 虚拟表配对。
virtual_table_eviction_policy (Optional[VirtualTableEvictionPolicy]) – 虚拟表的驱逐策略。

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

EmbeddingBagCollection 代表了一系列池化嵌入（EmbeddingBags）。

注意

EmbeddingBagCollection 是一个非分片的模块，没有针对性能进行优化。对于性能敏感的场景，请考虑使用分片版本 ShardedEmbeddingBagCollection。

它可以接收表示稀疏数据的参数，这些数据以 KeyedJaggedTensor 的形式存在，形状为 (F, B, L[f][i])，其中：

F：特征（键）的数量
B：批量大小
L[f][i]：稀疏特征的长度（可能对每个特征 f 和批量索引 i 不同，即锯齿状）

并输出一个 KeyedTensor，其值的形状为 (B, D)，其中：

B：批量大小
D：所有嵌入表的嵌入维度之和，即 sum([config.embedding_dim for config in tables])。

假设输入是一个 KeyedJaggedTensor J，包含 F 个特征、批量大小 B 和 L[f][i] 的稀疏长度，使得 J[f][i] 是特征 f 和批量索引 i 的 bag，则输出 KeyedTensor KT 定义如下：KT[i] = torch.cat([emb[f](J[f][i]) for f in J.keys()])，其中 emb[f] 是与特征 f 对应的 EmbeddingBag。

请注意，J[f][i] 是一个变长整数列表（一个 bag），而 emb[f](J[f][i]) 是通过使用 EmbeddingBag emb[f] 的模式（默认为平均值）来减少 J[f][i] 中每个值的嵌入而产生的池化嵌入。

参数:

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

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

features = KeyedJaggedTensor(
    keys=["f1", "f2"],
    values=torch.tensor([0, 1,                  2,    # feature 'f1'
                            3,      4,    5, 6, 7]),  # feature 'f2'
                    #    i = 1    i = 2    i = 3   <--- batch indices
    offsets=torch.tensor([
            0, 2, 2,       # 'f1' bags are values[0:2], values[2:2], and values[2:3]
            3, 4, 5, 8]),  # 'f2' bags are values[3:4], values[4:5], and values[5:8]
)

pooled_embeddings = ebc(features)
print(pooled_embeddings.values())
tensor([
    #  f1 pooled embeddings              f2 pooled embeddings
    #     from bags (dim. 3)                from bags (dim. 4)
    [-0.8899, -0.1342, -1.9060,  -0.0905, -0.2814, -0.9369, -0.7783],  # i = 0
    [ 0.0000,  0.0000,  0.0000,   0.1598,  0.0695,  1.3265, -0.1011],  # i = 1
    [-0.4256, -1.1846, -2.1648,  -1.0893,  0.3590, -1.9784, -0.7681]],  # i = 2
    grad_fn=<CatBackward0>)
print(pooled_embeddings.keys())
['f1', 'f2']
print(pooled_embeddings.offset_per_key())
tensor([0, 3, 7])  # embeddings have dimensions 3 and 4, so embeddings are at [0, 3) and [3, 7).

property device: device¶: 返回：torch.device：计算设备。

embedding_bag_configs() → List[EmbeddingBagConfig]¶

返回:: 嵌入 bag 配置。
返回类型:: List[EmbeddingBagConfig]

forward(features: KeyedJaggedTensor) → KeyedTensor¶

运行 EmbeddingBagCollection 的前向传播。此方法接收一个 KeyedJaggedTensor 并返回一个 KeyedTensor，这是对每个特征的嵌入进行池化后的结果。

参数:: features (KeyedJaggedTensor) – 输入 KJT
返回:: KeyedTensor

is_weighted() → bool¶

返回:: EmbeddingBagCollection 是否加权。
返回类型:: 布尔值

reset_parameters() → None¶: 重置 EmbeddingBagCollection 的参数。参数值根据每个 EmbeddingBagConfig 的 init_fn（如果存在）进行初始化。

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

EmbeddingCollection 代表了一系列非池化嵌入。

注意

EmbeddingCollection 是一个非分片的模块，没有针对性能进行优化。对于性能敏感的场景，请考虑使用分片版本 ShardedEmbeddingCollection。

它可以接收表示稀疏数据的参数，这些数据以 KeyedJaggedTensor 的形式存在，形状为 (F, B, L[f][i])，其中：

F：特征（键）的数量
B：批量大小
L[f][i]：稀疏特征的长度（可能对每个特征 f 和批量索引 i 不同，即锯齿状）

并输出一个 result，类型为 Dict[Feature, JaggedTensor]，其中 result[f] 是一个 JaggedTensor，形状为 (EB[f], D[f])，其中：

EB[f]：特征 f 的“扩展批量大小”，等于其 bag 值长度的总和，即 sum([len(J[f][i]) for i in range(B)])。
D[f]：特征 f 的嵌入维度。

参数:

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,    # feature 'f1'
                            3,      4,    5, 6, 7]),  # feature 'f2'
                    #    i = 1    i = 2    i = 3   <--- batch indices
    offsets=torch.tensor([
            0, 2, 2,       # 'f1' bags are values[0:2], values[2:2], and values[2:3]
            3, 4, 5, 8]),  # 'f2' bags are values[3:4], values[4:5], and values[5:8]
)

feature_embeddings = ec(features)
print(feature_embeddings['f2'].values())
tensor([
    # embedding for value 3 in f2 bag values[3:4]:
    [-0.2050,  0.5478,  0.6054],

    # embedding for value 4 in f2 bag values[4:5]:
    [ 0.7352,  0.3210, -3.0399],

    # embedding for values 5, 6, 7 in f2 bag values[5:8]:
    [ 0.1279, -0.1756, -0.4130],
    [ 0.7519, -0.4341, -0.0499],
    [ 0.9329, -1.0697, -0.8095],

], grad_fn=<EmbeddingBackward>)

property device: device¶: 返回：torch.device：计算设备。

embedding_configs() → List[EmbeddingConfig]¶

返回:: 嵌入配置。
返回类型:: List[EmbeddingConfig]

embedding_dim() → int¶

返回:: 嵌入维度。
返回类型:: int

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

返回:: 表中嵌入的名称。
返回类型:: List[List[str]]

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

运行 EmbeddingBagCollection 的前向传播。此方法接收一个 KeyedJaggedTensor 并返回一个 Dict[str, JaggedTensor]，这是对每个特征的单独嵌入的结果。

参数:: features (KeyedJaggedTensor) – KJT，形式为 [F X B X L]。
返回:: Dict[str, JaggedTensor]

need_indices() → bool¶

返回:: EmbeddingCollection 是否需要索引。
返回类型:: 布尔值

reset_parameters() → None¶: 重置 EmbeddingCollection 的参数。参数值根据每个 EmbeddingConfig 的 init_fn（如果存在）进行初始化。

模块¶

文档

教程

资源