1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
|
import os.path as osp
import torch
from tqdm import tqdm
from torch_geometric.datasets import AmazonBook
from torch_geometric.nn import LightGCN
from torch_geometric.utils import degree
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Amazon')
dataset = AmazonBook(path)
data = dataset[0]
num_users, num_books = data['user'].num_nodes, data['book'].num_nodes
data = data.to_homogeneous().to(device)
# Use all message passing edges as training labels:
batch_size = 8192
mask = data.edge_index[0] < data.edge_index[1]
train_edge_label_index = data.edge_index[:, mask]
train_loader = torch.utils.data.DataLoader(
range(train_edge_label_index.size(1)),
shuffle=True,
batch_size=batch_size,
)
model = LightGCN(
num_nodes=data.num_nodes,
embedding_dim=64,
num_layers=2,
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
def train():
total_loss = total_examples = 0
for index in tqdm(train_loader):
# Sample positive and negative labels.
pos_edge_label_index = train_edge_label_index[:, index]
neg_edge_label_index = torch.stack([
pos_edge_label_index[0],
torch.randint(num_users, num_users + num_books,
(index.numel(), ), device=device)
], dim=0)
edge_label_index = torch.cat([
pos_edge_label_index,
neg_edge_label_index,
], dim=1)
optimizer.zero_grad()
pos_rank, neg_rank = model(data.edge_index, edge_label_index).chunk(2)
loss = model.recommendation_loss(
pos_rank,
neg_rank,
node_id=edge_label_index.unique(),
)
loss.backward()
optimizer.step()
total_loss += float(loss) * pos_rank.numel()
total_examples += pos_rank.numel()
return total_loss / total_examples
@torch.no_grad()
def test(k: int):
emb = model.get_embedding(data.edge_index)
user_emb, book_emb = emb[:num_users], emb[num_users:]
precision = recall = total_examples = 0
for start in range(0, num_users, batch_size):
end = start + batch_size
logits = user_emb[start:end] @ book_emb.t()
# Exclude training edges:
mask = ((train_edge_label_index[0] >= start) &
(train_edge_label_index[0] < end))
logits[train_edge_label_index[0, mask] - start,
train_edge_label_index[1, mask] - num_users] = float('-inf')
# Computing precision and recall:
ground_truth = torch.zeros_like(logits, dtype=torch.bool)
mask = ((data.edge_label_index[0] >= start) &
(data.edge_label_index[0] < end))
ground_truth[data.edge_label_index[0, mask] - start,
data.edge_label_index[1, mask] - num_users] = True
node_count = degree(data.edge_label_index[0, mask] - start,
num_nodes=logits.size(0))
topk_index = logits.topk(k, dim=-1).indices
isin_mat = ground_truth.gather(1, topk_index)
precision += float((isin_mat.sum(dim=-1) / k).sum())
recall += float((isin_mat.sum(dim=-1) / node_count.clamp(1e-6)).sum())
total_examples += int((node_count > 0).sum())
return precision / total_examples, recall / total_examples
for epoch in range(1, 101):
loss = train()
precision, recall = test(k=20)
print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Precision@20: '
f'{precision:.4f}, Recall@20: {recall:.4f}')
|
