lightgcn.py

1"""2LightGCN 图神经网络召回（第10章 §10.2.4）34LightGCN 去掉了 GCN 中的特征变换和非线性激活，5仅保留图上的邻域聚合（消息传递），训练速度快、效果好。67参考论文：He et al., "LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation", SIGIR 202089消息传播公式：10    e_u^{(k+1)} = Σ_{i ∈ N(u)} e_i^{(k)} / sqrt(|N(u)| * |N(i)|)11    e_i^{(k+1)} = Σ_{u ∈ N(i)} e_u^{(k)} / sqrt(|N(i)| * |N(u)|)1213最终 Embedding = 各层平均1415实现说明：模型为 PyTorch nn.Module，归一化邻接矩阵用 torch tensor 存储，16    消息传递即稀疏矩阵乘法，梯度由 autograd 自动求导。17"""1819import math20import random21from collections import defaultdict2223import torch24import torch.nn as nn25import torch.optim as optim262728class LightGCN(nn.Module):29    def __init__(30        self,31        n_users: int,32        n_items: int,33        dim: int = 32,34        n_layers: int = 3,35        lr: float = 0.01,36        reg: float = 1e-4,37    ):38        super().__init__()39        self.n_users = n_users40        self.n_items = n_items41        self.dim = dim42        self.n_layers = n_layers43        self.lr = lr44        self.reg = reg4546        scale = 1.0 / math.sqrt(dim)47        self.user_emb0 = nn.Embedding(n_users, dim)48        self.item_emb0 = nn.Embedding(n_items, dim)49        nn.init.uniform_(self.user_emb0.weight, -scale, scale)50        nn.init.uniform_(self.item_emb0.weight, -scale, scale)5152        self.n_nodes = n_users + n_items53        self.adj_norm = torch.eye(self.n_nodes)  # Â = D^(-1/2) A D^(-1/2)5455    def build_graph(self, interactions: list[tuple[int, int]]) -> None:56        """构建对称归一化邻接矩阵（torch 稠密 tensor，本文件图规模小）。"""57        user_items: dict[int, list[int]] = defaultdict(list)58        item_users: dict[int, list[int]] = defaultdict(list)59        for u, i in interactions:60            user_items[u].append(i)61            item_users[i].append(u)6263        adj = torch.zeros(self.n_nodes, self.n_nodes)64        for u, items in user_items.items():65            du = len(items)66            for i in items:67                di = len(item_users[i])68                coeff = 1.0 / math.sqrt(du * di)69                adj[u, self.n_users + i] = coeff70                adj[self.n_users + i, u] = coeff71        self.adj_norm = adj7273    def propagate(self) -> tuple[torch.Tensor, torch.Tensor]:74        """75        LightGCN 前向传播：返回各层平均后的最终用户/物品 Embedding。76        e^(k+1) = Â e^(k)，最终 Embedding = 各层平均。77        """78        # 拼接为 (N, dim) 的统一节点表示：用户在前，物品在后79        h = torch.cat([self.user_emb0.weight, self.item_emb0.weight], dim=0)80        layers = [h]81        for _ in range(self.n_layers):82            h = self.adj_norm @ h   # 邻域聚合（已含度归一化）83            layers.append(h)8485        final = torch.stack(layers, dim=0).mean(dim=0)  # 各层平均86        final_u = final[: self.n_users]87        final_i = final[self.n_users:]88        return final_u, final_i8990    def train_model(91        self,92        interactions: list[tuple[int, int]],93        epochs: int = 10,94        batch_size: int = 16,95    ) -> None:96        """BPR 损失 + LightGCN 传播，梯度由 PyTorch autograd 自动求导。"""97        self.build_graph(interactions)98        user_pos: dict[int, set[int]] = defaultdict(set)99        for u, i in interactions:100            user_pos[u].add(i)101102        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.reg)103        samples = list(interactions)104105        for epoch in range(epochs):106            random.shuffle(samples)107            total_loss, n_batch = 0.0, 0108            for start in range(0, len(samples), batch_size):109                batch = samples[start:start + batch_size]110                u_ids, p_ids, n_ids = [], [], []111                for u, pos in batch:112                    neg = random.randint(0, self.n_items - 1)113                    while neg in user_pos[u]:114                        neg = random.randint(0, self.n_items - 1)115                    u_ids.append(u)116                    p_ids.append(pos)117                    n_ids.append(neg)118119                eu, ei = self.propagate()120                eu_b = eu[torch.tensor(u_ids, dtype=torch.long)]121                ei_p = ei[torch.tensor(p_ids, dtype=torch.long)]122                ei_n = ei[torch.tensor(n_ids, dtype=torch.long)]123                diff = (eu_b * ei_p).sum(-1) - (eu_b * ei_n).sum(-1)124                loss = -torch.log(torch.sigmoid(diff) + 1e-10).mean()125126                optimizer.zero_grad()127                loss.backward()128                optimizer.step()129                total_loss += loss.item()130                n_batch += 1131132            if epoch % 2 == 0:133                print(f"  epoch {epoch:3d}  avg_loss={total_loss / max(n_batch, 1):.4f}")134135    def recommend(136        self,137        uid: int,138        interactions: list[tuple[int, int]],139        top_k: int = 10,140    ) -> list[tuple[int, float]]:141        with torch.no_grad():142            eu, ei = self.propagate()143        user_items: dict[int, list[int]] = defaultdict(list)144        for u, i in interactions:145            user_items[u].append(i)146        seen = set(user_items.get(uid, []))147        scores = [148            (iid, (eu[uid] * ei[iid]).sum().item())149            for iid in range(self.n_items)150            if iid not in seen151        ]152        scores.sort(key=lambda x: x[1], reverse=True)153        return scores[:top_k]154155156# ── 快速验证 ──────────────────────────────────────────────────────────────────157if __name__ == "__main__":158    random.seed(42)159    torch.manual_seed(42)160    n_users, n_items = 8, 12161    interactions = list(set((random.randint(0, n_users - 1), random.randint(0, n_items - 1)) for _ in range(40)))162163    model = LightGCN(n_users, n_items, dim=8, n_layers=2, lr=0.05)164    print("LightGCN 训练：")165    model.train_model(interactions, epochs=6)166167    recs = model.recommend(0, interactions, top_k=5)168    print("\n用户 0 推荐结果：")169    for iid, s in recs:170        print(f"  item {iid}: {s:.4f}")