Source code for libauc.models.neumf

# This implementation is from

import torch
import torch.nn as nn
import logging

[docs] class NeuMF(nn.Module): r""" NeuMF is a widely-used model for recommender systems. args: user_num (int): the number of users in the dataset item_num (int): the number of items in the dataset dropout (float, optional): dropout ratio for the model emb_size (int, optional): embedding size of the model layers (string, optional): describe the layer information of the model Reference: .. [1] He, X., Liao, L., Zhang, H., Nie, L., Hu, X., and Chua, T. Neural Collaborative Filtering """ def __init__(self, user_num: int, item_num: int, dropout: float=0.2, emb_size: int=64, layers: str='[64]'): super(NeuMF, self).__init__() self.user_num = user_num self.item_num = item_num self.emb_size = emb_size self.dropout = dropout self.layers = eval(layers) self.mf_u_embeddings = nn.Embedding(self.user_num, self.emb_size) self.mf_i_embeddings = nn.Embedding(self.item_num, self.emb_size) self.mlp_u_embeddings = nn.Embedding(self.user_num, self.emb_size) self.mlp_i_embeddings = nn.Embedding(self.item_num, self.emb_size) self.mlp = nn.ModuleList([]) pre_size = 2 * self.emb_size for i, layer_size in enumerate(self.layers): self.mlp.append(nn.Linear(pre_size, layer_size)) pre_size = layer_size self.dropout_layer = nn.Dropout(p=self.dropout) self.prediction = nn.Linear(pre_size + self.emb_size, 1, bias=False)
[docs] def reset_last_layer(self): self.prediction.reset_parameters()
[docs] @staticmethod def init_weights(m): if 'Linear' in str(type(m)): nn.init.normal_(m.weight, mean=0.0, std=0.01) if m.bias is not None: nn.init.normal_(m.bias, mean=0.0, std=0.01) elif 'Embedding' in str(type(m)): nn.init.normal_(m.weight, mean=0.0, std=0.01)
[docs] def save_model(self, model_path=None): if model_path is None: model_path = self.model_path, model_path)
[docs] def load_model(self, model_path=None): if model_path is None: model_path = self.model_path self.load_state_dict(torch.load(model_path))'Load model from ' + model_path)
def forward(self, feed_dict): u_ids = feed_dict['user_id'].long() # [batch_size] i_ids = feed_dict['item_id'].long() # [batch_size, -1] u_ids = u_ids.unsqueeze(-1).repeat((1, i_ids.shape[1])) # [batch_size, -1] mf_u_vectors = self.mf_u_embeddings(u_ids) mf_i_vectors = self.mf_i_embeddings(i_ids) mlp_u_vectors = self.mlp_u_embeddings(u_ids) mlp_i_vectors = self.mlp_i_embeddings(i_ids) mf_vector = mf_u_vectors * mf_i_vectors mlp_vector =[mlp_u_vectors, mlp_i_vectors], dim=-1) for layer in self.mlp: mlp_vector = layer(mlp_vector).relu() mlp_vector = self.dropout_layer(mlp_vector) output_vector =[mf_vector, mlp_vector], dim=-1) prediction = self.prediction(output_vector) return {'prediction': prediction.view(feed_dict['batch_size'], -1)}