pinsage.py

1"""2PinSage 工业级图卷积召回（第10章 §10.2.2 图网络召回 —— PinSage召回）34核心思想：5    PinSage 没有发明新的图卷积结构，而是在 GraphSAGE 框架上做了面向超大规模推荐6    系统的工程化改造，三个最核心的创新点：7    1) Importance Sampling：用 Random Walk with Restart（RWR）游走，以访问频次8       （Visit Count）衡量邻居重要性，只取 Top-T 个重要邻居；9    2) Importance Pooling：按重要性权重做加权平均，而不是等权 mean/max；10    3) Margin Ranking Loss + Hard Negative：训练目标直接对齐检索排序。1112核心公式（对应正文公式）：13    访问频次 → α_u = count(u) / Σ_{u'∈N(v)} count(u')      # Importance Weight14    m_u = ReLU(Q h_u + q)                                   # 共享 MLP15    n_v = Σ_{u∈N(v)} α_u m_u                                # Importance Pooling（加权平均）16    z_v = ReLU(W [z_v ‖ n_v] + w)，z_v ← z_v / ||z_v||_2     # CONCAT + 线性 + L2 归一化17    L = Σ max(0, z_q·z_j - z_q·z_i + Δ)                      # Margin Ranking Loss1819与 GraphSAGE / LightGCN 的关键区别：20    - GraphSAGE 对邻居「随机均匀采样 + 等权聚合」，PinSage 用 RWR 访问频次做21      「重要性采样 + 重要性加权聚合」，更聚焦真实兴趣传播路径；22    - GraphSAGE 用 Skip-Gram 式无监督目标，PinSage 用 Margin Ranking Loss，23      并刻意构造 Hard Negative（有点像但最终没被点击的物品）；24    - LightGCN 要全图传播，PinSage 只在 Mini-batch 的 K 跳子图上卷积，25      访问频次权重还能由 CPU 离线预计算，天然适合十亿级图。2627参考论文：Ying et al., "Graph Convolutional Neural Networks for Web-Scale28         Recommender Systems", KDD 20182930实现说明：模型为 PyTorch nn.Module，邻接表用 Python list 存储，消息传递与线性变换用31    tensor 操作完成，梯度由 autograd 自动求导。32"""3334import math35import random36from collections import defaultdict3738import torch39import torch.nn as nn40import torch.nn.functional as F41import torch.optim as optim424344# ── 基础工具函数（本文件自包含，不跨文件 import） ────────────────────────────45def _cosine(a: list[float], b: list[float]) -> float:46    na, nb = math.sqrt(sum(x * x for x in a)), math.sqrt(sum(x * x for x in b))47    if na < 1e-12 or nb < 1e-12:48        return 0.049    return sum(x * y for x, y in zip(a, b)) / (na * nb)505152class PinSageRecall(nn.Module):53    """基于 RWR 重要性采样 + Importance Pooling 的 PinSage 召回模型。54    PyTorch autograd 自动求导。"""5556    def __init__(57        self, n_users: int, n_items: int, dim: int = 16, n_layers: int = 2, top_t: int = 5,58        margin: float = 0.2, lr: float = 0.05, reg: float = 1e-4,59    ):60        super().__init__()61        self.n_users, self.n_items = n_users, n_items62        self.n_nodes = n_users + n_items  # 用户 u -> u，物品 i -> n_users + i63        self.dim, self.n_layers, self.top_t = dim, n_layers, top_t64        self.margin, self.lr, self.reg = margin, lr, reg6566        scale = 1.0 / math.sqrt(dim)67        # 初始特征取非负（模拟 PinSage 中经过 ReLU 的视觉/文本特征），避免逐层 ReLU 把信号截断为 068        self.emb = nn.Embedding(self.n_nodes, dim)69        nn.init.uniform_(self.emb.weight, 0.0, 2 * scale)7071        # Q, q：邻居共享的 MLP (dim -> dim)；W, w：CONCAT (2*dim -> dim)72        self.Q = nn.ModuleList([nn.Linear(dim, dim) for _ in range(n_layers)])73        self.W = nn.ModuleList([nn.Linear(2 * dim, dim) for _ in range(n_layers)])74        for lin in self.Q + self.W:75            nn.init.xavier_normal_(lin.weight)76            nn.init.zeros_(lin.bias)7778        self.adj: list[list[int]] = []79        self.importance: list[list[tuple[int, float]]] = []  # 每个节点的 Top-T 重要邻居及权重80        self.visit_count: list[list[tuple[int, int]]] = []   # 原始访问频次（便于观察）81        self.node_emb: torch.Tensor | None = None8283    # ── 图构建 + RWR 重要性采样 ───────────────────────────────────────────84    def build_graph(self, interactions: list[tuple[int, int]]) -> None:85        """构建二部图邻接表（PinSage 工程上即 Pin-Board 图）。"""86        adj: dict[int, set[int]] = defaultdict(set)87        for u, i in interactions:88            adj[u].add(self.n_users + i)89            adj[self.n_users + i].add(u)90        self.adj = [sorted(adj[v]) for v in range(self.n_nodes)]9192    def sample_importance_neighbors(self, n_walks: int = 60, walk_len: int = 6, restart: float = 0.5) -> None:93        """RWR：每步以 restart 概率跳回起点；访问频次作为邻居重要性，取 Top-T 并归一化。"""94        self.importance, self.visit_count = [], []95        for v in range(self.n_nodes):96            counts: dict[int, int] = defaultdict(int)97            for _ in range(n_walks):98                cur = v99                for _ in range(walk_len):100                    if random.random() < restart:101                        cur = v102                        continue103                    if not self.adj[cur]:104                        break105                    cur = random.choice(self.adj[cur])106                    # 只统计物品节点：用户节点得到 U2I 邻居，物品节点得到二跳的 I2I 邻居107                    if cur != v and cur >= self.n_users:108                        counts[cur] += 1109            top = sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))[: self.top_t]110            total = sum(c for _, c in top) or 1  # α_u = count(u) / Σ count111            self.importance.append([(n, c / total) for n, c in top])112            self.visit_count.append(top)113114    # ── 前向传播（PinSage Convolution 模块） ──────────────────────────────115    def propagate(self) -> torch.Tensor:116        """逐层执行 Importance Pooling + CONCAT + 线性 + L2 归一化，返回 (N, dim)。"""117        h = self.emb.weight   # (N, dim)118        for layer in range(self.n_layers):119            new_h_list = []120            for v in range(self.n_nodes):121                # (1) 共享 MLP：m_u = ReLU(Q h_u + q)；(2) Importance Pooling：n_v = Σ α_u m_u122                if self.importance[v]:123                    nbr_ids = [u for u, _ in self.importance[v]]124                    alphas_t = torch.tensor([a for _, a in self.importance[v]], dtype=torch.float)125                    h_nbrs = h[torch.tensor(nbr_ids, dtype=torch.long)]   # (T, dim)126                    m = torch.relu(self.Q[layer](h_nbrs))                 # (T, dim)127                    n_v = (alphas_t.unsqueeze(-1) * m).sum(dim=0)         # (dim,)128                else:129                    n_v = torch.zeros(self.dim)130                # (3) CONCAT(自身, 邻域) 后线性变换 + ReLU + L2 归一化131                x = torch.cat([h[v], n_v], dim=0)    # (2*dim,)132                pre = torch.relu(self.W[layer](x))    # (dim,)133                new_h_list.append(F.normalize(pre, p=2, dim=0))134            h = torch.stack(new_h_list, dim=0)135        self.node_emb = h136        return h137138    # ── Hard Negative 采样 ────────────────────────────────────────────────139    def _hard_negative(self, q_node: int, pos_items: set[int], emb: torch.Tensor, pool: int = 5) -> int:140        """从随机候选池里挑与 query 最相似但用户没点过的物品，构造 Hard Negative。"""141        cands = [j for j in (random.randrange(self.n_items) for _ in range(pool)) if j not in pos_items]142        if not cands:143            return random.randrange(self.n_items)144        eu = emb[q_node]145        best = max(cands, key=lambda j: (eu * emb[self.n_users + j]).sum().item())146        return best147148    # ── 训练 ──────────────────────────────────────────────────────────────149    def train_model(self, interactions: list[tuple[int, int]], epochs: int = 6) -> None:150        """Margin Ranking Loss 训练，梯度由 PyTorch autograd 自动求导。"""151        self.build_graph(interactions)152        self.sample_importance_neighbors()153        user_pos: dict[int, set[int]] = defaultdict(set)154        for u, i in interactions:155            user_pos[u].add(i)156157        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.reg)158        samples = list(interactions)159160        for epoch in range(epochs):161            random.shuffle(samples)162            total_loss, n_violate = 0.0, 0163164            # Full forward once per epoch (recompute per sample is too slow)165            h = self.propagate()166            batch_u, batch_p, batch_n = [], [], []167            for u, pos in samples:168                if random.random() < 0.5:169                    neg = self._hard_negative(u, user_pos[u], h.detach())170                else:171                    neg = random.randrange(self.n_items)172                    while neg in user_pos[u]:173                        neg = random.randrange(self.n_items)174                batch_u.append(u)175                batch_p.append(self.n_users + pos)176                batch_n.append(self.n_users + neg)177178            h = self.propagate()  # recompute with grad179            z_u = h[torch.tensor(batch_u, dtype=torch.long)]180            z_p = h[torch.tensor(batch_p, dtype=torch.long)]181            z_n = h[torch.tensor(batch_n, dtype=torch.long)]182            # L = max(0, z_q·z_j - z_q·z_i + Δ)183            raw_loss = (z_u * z_n).sum(-1) - (z_u * z_p).sum(-1) + self.margin184            loss = torch.clamp(raw_loss, min=0.0).mean()185186            optimizer.zero_grad()187            loss.backward()188            optimizer.step()189190            total_loss = loss.item()191            n_violate = (raw_loss > 0).sum().item()192            print(f"  epoch {epoch:2d}  margin_loss={total_loss:.4f}  violate={n_violate}/{len(samples)}")193194    # ── 召回 ──────────────────────────────────────────────────────────────195    def recommend(self, uid: int, interactions: list[tuple[int, int]], top_k: int = 5):196        """U2I 召回：用户向量检索 Top-K 物品。"""197        with torch.no_grad():198            emb = self.propagate()199        seen = {i for u, i in interactions if u == uid}200        z_u = emb[uid].tolist()201        scores = [(i, _cosine(z_u, emb[self.n_users + i].tolist()))202                  for i in range(self.n_items) if i not in seen]203        scores.sort(key=lambda x: x[1], reverse=True)204        return scores[:top_k]205206    def similar_items(self, iid: int, top_k: int = 5):207        """I2I 召回：PinSage 线上最常用的形态（只保留物品 Embedding 建 ANN 索引）。"""208        with torch.no_grad():209            emb = self.propagate()210        base = emb[self.n_users + iid].tolist()211        scores = [(j, _cosine(base, emb[self.n_users + j].tolist()))212                  for j in range(self.n_items) if j != iid]213        scores.sort(key=lambda x: x[1], reverse=True)214        return scores[:top_k]215216217# ── 快速验证 ──────────────────────────────────────────────────────────────────218if __name__ == "__main__":219    random.seed(42)220    torch.manual_seed(42)221    n_users, n_items = 8, 20222223    # 构造带兴趣簇的合成二部图：前 4 个用户偏好物品 0~9，后 4 个偏好 10~19224    interactions = set()225    for u in range(n_users):226        lo, hi = (0, 9) if u < 4 else (10, 19)227        for _ in range(6):228            interactions.add((u, random.randint(lo, hi)))229        interactions.add((u, random.randrange(n_items)))  # 少量跨簇噪声230    interactions = sorted(interactions)231    print(f"合成二部图：{n_users} 个用户，{n_items} 个物品，{len(interactions)} 条交互边")232233    model = PinSageRecall(n_users, n_items, dim=16, n_layers=2, top_t=5, margin=0.2)234    print("\nPinSage 训练（RWR 重要性采样 + Importance Pooling + Margin Loss）：")235    model.train_model(interactions, epochs=6)236237    print("\n用户 0 的 Top-5 重要邻居（RWR 访问频次 → Importance Weight）：")238    for (node, cnt), (_, alpha) in zip(model.visit_count[0], model.importance[0]):239        print(f"  item {node - n_users:2d}  visit_count={cnt:3d}  α={alpha:.4f}")240241    for uid, desc in ((0, "偏好物品 0~9"), (7, "偏好物品 10~19")):242        print(f"\n用户 {uid}（{desc}）的 Top-5 U2I 召回：")243        for iid, s in model.recommend(uid, interactions, top_k=5):244            print(f"  item {iid:2d}  cosine={s:.4f}")245246    print("\n物品 3 的 Top-5 I2I 召回（线上以物品 Embedding 建 ANN 索引）：")247    for jid, s in model.similar_items(3, top_k=5):248        print(f"  item {jid:2d}  cosine={s:.4f}")