pytorch 实现FM算法

FM（因子分解机）是一种经典的推荐算法，它可以用于处理稀疏数据并且具有较好的预测性能。下面是使用PyTorch实现FM算法的基本步骤：

1. 导入需要的库：

import torch
import torch.nn as nn
import torch.optim as optim

2. 定义FM模型

class FM(nn.Module):
    def __init__(self, input_dim, k):
        super(FM, self).__init__()
        self.k = k
        self.linear = nn.Linear(input_dim, 1)
        self.v = nn.Parameter(torch.randn(input_dim, k))

    def forward(self, x):
        linear_part = self.linear(x)
        inter_part1 = torch.matmul(x, self.v)
        inter_part2 = torch.matmul(torch.pow(x, 2), torch.pow(self.v, 2))
        inter_part = 0.5 * torch.sum(torch.sub(inter_part1, inter_part2), 1, keepdim=True)
        output = linear_part + inter_part
        return output

3. 定义训练函数

def train(model, dataloader, optimizer, criterion):
    model.train()
    train_loss = 0
    for batch_idx, (data, target) in enumerate(dataloader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        train_loss += loss.item()
        loss.backward()
        optimizer.step()
    return train_loss / len(dataloader.dataset)

4. 定义测试函数

def test(model, dataloader, criterion):
    model.eval()
    test_loss = 0
    with torch.no_grad():
        for data, target in dataloader:
            output = model(data)
            test_loss += criterion(output, target).item()
    return test_loss / len(dataloader.dataset)

5. 加载数据集和设置超参数

from torch.utils.data import DataLoader, Dataset

class CustomDataset(Dataset):
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __getitem__(self, index):
        return self.x[index], self.y[index]

    def __len__(self):
        return len(self.x)

X_train, y_train = ...
X_test, y_test = ...
train_dataset = CustomDataset(X_train, y_train)
test_dataset = CustomDataset(X_test, y_test)
batch_size = 64
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
input_dim = X_train.shape[1]
k = 10
lr = 0.01
num_epochs = 50

6. 训练模型

model = FM(input_dim, k)
optimizer = optim.SGD(model.parameters(), lr=lr)
criterion = nn.MSELoss()
for epoch in range(num_epochs):
    train_loss = train(model, train_dataloader, optimizer, criterion)
    test_loss = test(model, test_dataloader, criterion)
    print(f'Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}')

这样就可以使用PyTorch实现FM算法了。