pytorch 推荐FM算法实现

Factorization Machine (FM) 是一种处理稀疏高维数据的算法，可以用于推荐系统、CTR 预估等场景。下面是使用 PyTorch 实现 FM 算法的示例代码。

首先，我们需要导入 PyTorch 和其他必要的库：

import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import LabelEncoder

然后我们定义一个自定义数据集类来加载数据：

class FMDataSet(Dataset):
    def __init__(self, X, y=None):
        self.X = X
        self.y = y
        
    def __len__(self):
        return len(self.X)
    
    def __getitem__(self, idx):
        if self.y is not None:
            return self.X[idx], self.y[idx]
        else:
            return self.X[idx]

接着，我们定义 FM 模型的类：

class FM(torch.nn.Module):
    def __init__(self, num_features, k):
        super(FM, self).__init__()
        self.num_features = num_features
        self.k = k
        # 线性部分
        self.linear = torch.nn.Linear(num_features, 1)
        # 交叉部分
        self.V = torch.nn.Parameter(torch.randn(num_features, k))
    
    def forward(self, x):
        linear_part = self.linear(x)
        v_part = torch.mm(x, self.V).pow(2).sum(1, keepdim=True) \
                 - torch.mm(x.pow(2), self.V.pow(2)).sum(1, keepdim=True)
        fm_part = 0.5 * v_part
        out = linear_part + fm_part
        return torch.sigmoid(out)

在这个模型中，我们定义了线性部分和交叉部分。线性部分使用 PyTorch 自带的 Linear 层，而交叉部分使用了一个矩阵 V，它的形状是 (num_features, k)。我们通过矩阵乘法和张量乘法计算出交叉部分的输出。

接下来，我们可以加载数据并进行预处理。这里我们使用 MovieLens 数据集作为示例数据：

import pandas as pd

# 加载数据
data = pd.read_csv('http://files.grouplens.org/datasets/movielens/ml-100k/u.data',
                   sep='\t', header=None, names=['user_id', 'item_id', 'rating', 'timestamp'])

# 将 user_id 和 item_id 编码成连续的整数
user_encoder = LabelEncoder()
item_encoder = LabelEncoder()
data['user_id'] = user_encoder.fit_transform(data['user_id'])
data['item_id'] = item_encoder.fit_transform(data['item_id'])

# 将数据集拆分成训练集和测试集
train_data = FMDataSet(data[['user_id', 'item_id']].values, data['rating'].values)
train_loader = DataLoader(train_data, batch_size=256, shuffle=True)

num_users = len(user_encoder.classes_)
num_items = len(item_encoder.classes_)

最后，我们可以训练模型并进行评估：

# 定义模型和优化器
model = FM(num_users + num_items, 20)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(10):
    for x, y in train_loader:
        optimizer.zero_grad()
        y_pred = model(x)
        loss = torch.nn.functional.binary_cross_entropy(y_pred, y.float().unsqueeze(1))
        loss.backward()
        optimizer.step()
    print('Epoch {}: loss={:.4f}'.format(epoch+1, loss.item()))

# 在测试集上评估模型
test_data = FMDataSet(data[['user_id', 'item_id']].values)
test_loader = DataLoader(test_data, batch_size=1024)
y_pred = []
for x in test_loader:
    y_pred.append(model(x).detach().numpy().flatten())
y_pred = np.concatenate(y_pred)
y_true = data['rating'].values
rmse = np.sqrt(np.mean((y_true - y_pred)**2))
print('RMSE={:.4f}'.format(rmse))

这样就可以使用 PyTorch 实现 FM 算法了。请注意，这个示例代码只是一个简单的示例，实际的应用中可能需要进行更多的优化和调整。