Python PyTorch SparseArch用法及代码示例

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

用法:

class torchrec.models.dlrm.SparseArch(embedding_bag_collection: torchrec.modules.embedding_modules.EmbeddingBagCollection)

参数：

embedding_bag_collection(torchrec.modules.embedding_modules.EmbeddingBagCollection) -表示池化嵌入的集合

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

处理DLRM的稀疏特征。对所有 EmbeddingBag 进行嵌入查找并嵌入每个集合的特征。

例子：

eb1_config = EmbeddingBagConfig(
   name="t1", embedding_dim=3, num_embeddings=10, feature_names=["f1"]
)
eb2_config = EmbeddingBagConfig(
   name="t2", embedding_dim=4, num_embeddings=10, feature_names=["f2"]
)
ebc_config = EmbeddingBagCollectionConfig(tables=[eb1_config, eb2_config])

ebc = EmbeddingBagCollection(config=ebc_config)
sparse_arch = SparseArch(embedding_bag_collection)

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

sparse_embedded = sparse_arch(features)