Python PyTorch FactorizationMachine用法及代碼示例

本文簡要介紹python語言中 torchrec.modules.deepfm.FactorizationMachine 的用法。

用法:

class torchrec.modules.deepfm.FactorizationMachine

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

這是 DeepFM paper 中提到的分解機模塊：

本模塊不涵蓋已發表論文的end-end 函數。相反，它僅涵蓋出版物的 FM 部分，並用於學習 2nd-order 特征交互。

為了支持建模靈活性，我們將關鍵組件定製為：

• 與公開論文不同，我們將輸入從原始稀疏

  特征到特征的嵌入。它允許嵌入維度和嵌入數量的靈活性，隻要所有嵌入張量具有相同的批量大小。

模塊的一般架構如下：

# 1 x 1 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)