Common neural network architectures for *Collaborative Filtering*.

Overview

This package implements several neural network architectures that can be used to build recommendation systems. Users of the library can add or define their own implementations or use the existing ones. There are two layers that every architecture should define:

  • user_embeddings: The user embedding matrix
  • item_embeddings: The item embedding matrix

Every implementation should be a subclass of torch.nn.Module.

Simple Collaborative Filtering

This architecture is the simplest one to implement Collaborative Filtering. It only defines the embedding matrices for users and items and the final rating is computed by the dot product of the corresponding rows.

class SimpleCF[source]

SimpleCF(n_users:int, n_items:int, factors:int=16, user_embeddings:tensor=None, freeze_users:bool=False, item_embeddings:tensor=None, freeze_items:bool=False, init:torch.nn.init='normal_', binary:bool=False, **kwargs) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

Arguments:

  • n_users (int): The number of unique users
  • n_items (int): The number of unique items
  • factors (int): The dimension of the embedding space
  • user_embeddings (torch.tensor): Pre-trained weights for the user embedding matrix
  • freeze_users (bool): True if we want to keep the user weights as is (i.e. non-trainable)
  • item_embeddings (torch.tensor): Pre-trained weights for the item embedding matrix
  • freeze_item (bool): True if we want to keep the item weights as is (i.e. non-trainable)
  • init (torch.nn.init): The initialization method of the embedding matrices - default: torch.nn.init.normal_
# initialize the model with 100 users, 50 items and a 16-dimensional embedding space
model = SimpleCF(100, 50, 16, mean=0., std=.1, binary=True)

# predict the rating that user 3 would give to item 33
model(torch.tensor([2]), torch.tensor([32]))
tensor([[[0.5009]]], grad_fn=<SigmoidBackward>)