推荐系统经典算法：FM的pytorch实现

最新推荐文章于 2025-11-24 02:54:01 发布

原创

最新推荐文章于 2025-11-24 02:54:01 发布 · 3.6k 阅读

16 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #算法 #推荐系统 #深度学习

首先说明为什么用pytorch实现该算法：因为暑假里尝试过在cpu里运行的简陋版本，运行速度过于缓慢；pytorch中有现成的embedding方法可以使用并且做梯度下降比较容易。

下面先说核心算法部分：

这里的时间复杂度仍为o（kn^2），下面是化简到o（kn）

具体实现：

class FactorizationMachine(nn.Module):

    def __init__(self, field_dims, embed_dim=4):
        super(FactorizationMachine, self).__init__()

        self.embed1 = FeaturesEmbedding(field_dims, 1)
        self.embed2 = FeaturesEmbedding(field_dims, embed_dim)
        self.bias = nn.Parameter(torch.zeros((1,)))

    def forward(self, x):
        # x shape: (batch_size, num_fields)
        # embed(x) shape: (batch_size, num_fields, embed_dim)
        square_sum = self.embed2(x).sum(dim=1).pow(2).sum(dim=1)
        sum_square = self.embed2(x).pow(2).sum(dim=1).sum(dim=1)

        output = self.embed1(x).squeeze(-1).sum(dim=1) + self.bias + (square_sum - sum_square) / 2
        output = torch.sigmoid(output).unsqueeze(-1)
        return output

最后附上全部代码：

FM.py

from utils import create_dataset, Trainer
from layer import Embedding, FeaturesEmbedding, EmbeddingsInteraction, MultiLayerPerceptron

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

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Training on [{}].'.format(device))

dataset = create_dataset(sample_num=100000, device=device)
field_dims, (train_X, train_y), (test_X, test_y) = dataset.train_valid_test_split()


class FactorizationMachine(nn.Module):

    def __init__(self, field_dims, embed_dim=4):
        super(FactorizationMachine, self).__init__()

        self.embed1 = FeaturesEmbedding(field_dims, 1)
        self.embed2 = FeaturesEmbedding(field_dims, embed_dim)
        self.bias = nn.Parameter(torch.zeros((1,)))

    def forward(self, x):
        # x shape: (batch_size, num_fields)
        # embed(x) shape: (batch_size, num_fields, embed_dim)
        square_sum = self.embed2(x).sum(dim=1).pow(2).sum(dim=1)
        sum_square = self.embed2(x).pow(2).sum(dim=1).sum(dim=1)

        output = self.embed1(x).s