利用pytorch实现 NFM 算法 并且使用auc进行评估

好的，下面是基于 PyTorch 实现 NFM 算法的代码，同时使用 AUC 进行评估：

首先，我们需要导入必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import roc_auc_score

接着，我们定义 NFM 模型的代码如下：

class NFM(nn.Module):
    def __init__(self, num_feats, embed_dim, hidden_dim):
        super(NFM, self).__init__()
        
        self.embedding_layer = nn.Embedding(num_feats, embed_dim)
        self.linear_layer = nn.Linear(embed_dim, hidden_dim)
        self.interaction_layer = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        embeddings = self.embedding_layer(x)
        square_sum = torch.sum(embeddings, dim=1) ** 2
        sum_square = torch.sum(embeddings ** 2, dim=1)
        
        linear = self.linear_layer(torch.mean(embeddings, dim=1))
        interaction = torch.sum(self.interaction_layer(embeddings), dim=1)
        
        output = linear + interaction - square_sum + sum_square
        
        return output

在这里，我们将 NFM 模型分为三个部分：

1. `embedding_layer`：嵌入层，将输入的特征进行嵌入，得到每个特征的向量表示。
2. `linear_layer`：线性层，对特征嵌入后得到的向量进行平均，得到一个一维向量，并通过一个全连接层进行线性变换。
3. `interaction_layer`：交互层，对每个特征嵌入后得到的向量进行交互计算，并通过一系列层进行非线性变换，最后得到一个交互权重。

在 `forward` 函数中，我们首先计算了平方和以及和的平方，然后分别对线性部分和交互部分进行计算，最后将线性部分、交互部分、平方和以及和的平方相加得到最终的输出。

接下来，我们定义训练函数和测试函数：

def train(model, train_loader, optimizer, criterion):
    model.train()
    
    for batch_idx, (X, y) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(X)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()
        
def test(model, test_loader):
    model.eval()
    
    y_true = []
    y_pred = []
    
    with torch.no_grad():
        for batch_idx, (X, y) in enumerate(test_loader):
            output = model(X)
            y_true.extend(y.tolist())
            y_pred.extend(output.tolist())
            
    auc = roc_auc_score(y_true, y_pred)
    print('AUC:', auc)

在这里，我们定义了一个训练函数和一个测试函数。训练函数中，我们首先将模型设为训练模式，然后对于每一个 batch 进行一次前向传播、计算 loss、反向传播和更新参数的操作。测试函数中，我们将模型设为评估模式，然后对于测试集中的每一个样本进行一次前向传播，并记录真实标签和预测概率，最后计算 AUC。

最后，我们读取数据、初始化模型、定义损失函数和优化器，并开始训练和测试：

# 读取数据
train_data = torch.Tensor([[1, 2, 3, 0], [4, 5, 6, 1], [7, 8, 9, 0], [10, 11, 12, 1]])
test_data = torch.Tensor([[1, 2, 3, 1], [4, 5, 6, 0], [7, 8, 9, 1], [10, 11, 12, 0]])

train_loader = torch.utils.data.DataLoader(train_data, batch_size=2, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=2, shuffle=True)

# 初始化模型
model = NFM(num_feats=4, embed_dim=2, hidden_dim=4)

# 定义损失函数和优化器
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练和测试
for epoch in range(10):
    train(model, train_loader, optimizer, criterion)
    test(model, test_loader)

在这里，我们使用了一个简单的数据集作为示例，共有 4 个样本，每个样本都有 4 个特征，其中前三个特征是类别型特征，最后一个特征是二分类标签。在初始化模型时，我们将类别型特征嵌入到 2 维向量中，然后将线性层和交互层的隐藏维度都设置为 4。在优化器中，我们使用了 Adam 算法作为优化器，学习率为 0.01。在训练过程中，我们进行了 10 轮迭代，每轮迭代后都对模型进行一次测试，并输出 AUC。