class NeuMF(nn.Module):
def __init__(self, num_factors, num_users, num_items, nums_hiddens):
super().__init__()
self.P = nn.Embedding(num_users, num_factors)
self.Q = nn.Embedding(num_items, num_factors)
self.U = nn.Embedding(num_users, num_factors)
self.V = nn.Embedding(num_items, num_factors)
mlp_layers = []
input_size = num_factors * 2
for num_hiddens in nums_hiddens:
mlp_layers.append(nn.Linear(input_size, num_hiddens))
mlp_layers.append(nn.ReLU())
input_size = num_hiddens
self.mlp = nn.Sequential(*mlp_layers)
# Output raw logits; BPRLoss applies sigmoid internally, so adding
# a sigmoid here would compose with it and squash gradients.
self.prediction_layer = nn.Linear(num_factors + nums_hiddens[-1], 1,
bias=False)
def forward(self, user_id, item_id):
p_mf = self.P(user_id)
q_mf = self.Q(item_id)
gmf = p_mf * q_mf
p_mlp = self.U(user_id)
q_mlp = self.V(item_id)
mlp = self.mlp(torch.cat([p_mlp, q_mlp], dim=1))
con_res = torch.cat([gmf, mlp], dim=1)
return self.prediction_layer(con_res)Implicit-feedback dataset
Model architecture
Two embedding tables per side (one for GMF, one for MLP); elementwise product on one side, concat→MLP on the other; final concat → linear → sigmoid score:
NeuMF architecture: GMF and MLP pathways are fused before scoring a user-item pair.
Each training instance: a (user, positive item) plus sampled negatives. The dataset class handles negative sampling on the fly:
class PRDataset(torch.utils.data.Dataset):
def __init__(self, users, items, candidates, num_items, test_items=None):
self.users = users
self.items = items
# Precompute each user's negative pool once: items the user has not
# interacted with in train AND not held out as a test positive
# (excluding test positives prevents leakage into the BPR loss).
all_items = set(range(num_items))
test_items = test_items or {}
self.neg_pool = {
u: list(all_items - set(candidates.get(u, [])) - set(test_items.get(u, [])))
for u in candidates}
def __len__(self):
return len(self.users)
def __getitem__(self, idx):
neg_items = self.neg_pool[int(self.users[idx])]
indices = random.randint(0, len(neg_items) - 1)
return self.users[idx], self.items[idx], neg_items[indices]Training and Metrics
The final printout should be read as ranking quality: higher Hit@50 and AUC mean the held-out item is placed above more unobserved candidates. They are not rating-prediction metrics:
lr, num_epochs, wd = 0.01, 10, 1e-5
loss = d2l.BPRLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd)
train_ranking(net, train_iter, test_iter, loss, optimizer, None, num_users,
num_items, num_epochs, devices, evaluate_ranking, candidates)
train loss 0.154, test hit rate 0.361, test AUC 0.881
17.2 examples/sec on [device(type='cuda', index=0)]