Python PyTorch DLRM用法及代码示例

本文简要介绍python语言中 torchrec.models.dlrm.DLRM 的用法。

用法:

class torchrec.models.dlrm.DLRM(embedding_bag_collection: torchrec.modules.embedding_modules.EmbeddingBagCollection, dense_in_features: int, dense_arch_layer_sizes: List[int], over_arch_layer_sizes: List[int], dense_device: Optional[torch.device] = None)

参数：

• embedding_bag_collection(torchrec.modules.embedding_modules.EmbeddingBagCollection) -用于定义 SparseArch 的嵌入包集合。

• dense_in_features(int) -密集输入特征的维度。

• dense_arch_layer_sizes(List[int]) -DenseArch 的层大小。

• over_arch_layer_sizes(List[int]) -OverArch 的层大小。 InteractionArch 的输出维度不应在此处手动指定。

• dense_device(可选的[torch.device]) -默认计算设备。

基础：torch.nn.modules.module.Module

Recsys 模型来自“个性化和推荐系统的深度学习推荐模型”(https://arxiv.org/abs/1906.00091)。通过学习每个特征的池化嵌入来处理稀疏特征。通过将密集特征投影到相同的嵌入空间来学习密集特征和稀疏特征之间的关系。此外，还学习稀疏特征之间的成对关系。

该模块假设所有稀疏特征具有相同的嵌入维度(即每个EmbeddingBagConfig使用相同的embedding_dim)。

在整个模型的文档中使用以下符号：

• F：稀疏特征的数量

• D：embedding_dimension 稀疏特征

• B：批量大小

• num_features：密集特征的数量

例子：

B = 2
D = 8

eb1_config = EmbeddingBagConfig(
   name="t1", embedding_dim=D, num_embeddings=100, feature_names=["f1", "f3"]
)
eb2_config = EmbeddingBagConfig(
   name="t2",
   embedding_dim=D,
   num_embeddings=100,
   feature_names=["f2"],
)
ebc_config = EmbeddingBagCollectionConfig(tables=[eb1_config, eb2_config])

ebc = EmbeddingBagCollection(config=ebc_config)
model = DLRM(
   embedding_bag_collection=ebc,
   dense_in_features=100,
   dense_arch_layer_sizes=[20],
   over_arch_layer_sizes=[5, 1],
)

features = torch.rand((B, 100))

#     0       1
# 0   [1,2] [4,5]
# 1   [4,3] [2,9]
# ^
# feature
sparse_features = KeyedJaggedTensor.from_offsets_sync(
   keys=["f1", "f3"],
   values=torch.tensor([1, 2, 4, 5, 4, 3, 2, 9]),
   offsets=torch.tensor([0, 2, 4, 6, 8]),
)

logits = model(
   dense_features=features,
   sparse_features=sparse_features,
)