sasrec.py

1"""2SASRec Transformer 序列召回（第10章 §10.4.3）34论文：Kang & McAuley, "Self-Attentive Sequential Recommendation", ICDM 201856核心思想：7- 用单向（Causal）Self-Attention 对用户行为序列建模8- 每个位置的输出作为「下一个物品」的预测向量9- 训练用 BCE loss + 负采样；线上取最后一个位置的输出做 ANN 检索1011本实现：PyTorch autograd，使用 nn.TransformerEncoderLayer 配合 causal mask，12完整 BPTT 由 autograd 自动完成。13"""1415import math16import random17from collections import defaultdict1819import torch20import torch.nn as nn21import torch.nn.functional as F22import torch.optim as optim232425# ── SASRec 模型 ───────────────────────────────────────────────────────────────2627class SASRec(nn.Module):28    """29    单层、单头 SASRec。3031    组件：32      item_emb: nn.Embedding [n_items × dim]33      pos_emb:  nn.Embedding [max_len × dim]34      transformer_layer: nn.TransformerEncoderLayer (causal mask)35      输出 item Embedding 与 item_emb 共享（Tied Embedding 风格）36    """3738    def __init__(39        self,40        n_items: int,41        dim: int = 32,42        max_len: int = 50,43        lr: float = 0.01,44        reg: float = 1e-4,45        n_heads: int = 1,46        dropout: float = 0.1,47    ):48        super().__init__()49        self.n_items = n_items50        self.dim = dim51        self.max_len = max_len52        self.lr = lr53        self.reg = reg5455        scale = math.sqrt(1.0 / dim)56        self.item_emb = nn.Embedding(n_items, dim)57        self.pos_emb = nn.Embedding(max_len, dim)58        nn.init.uniform_(self.item_emb.weight, -scale, scale)59        nn.init.uniform_(self.pos_emb.weight, -scale, scale)6061        # 单层因果 Transformer Encoder（batch_first=True）62        self.transformer_layer = nn.TransformerEncoderLayer(63            d_model=dim,64            nhead=n_heads,65            dim_feedforward=dim * 4,66            dropout=dropout,67            activation="relu",68            batch_first=True,69            norm_first=False,70        )7172    def _encode(self, seq_ids: torch.Tensor) -> torch.Tensor:73        """74        seq_ids: [B, T]  item id 序列（已截断到 max_len）75        返回: [B, T, dim]  每个位置的上下文表示76        """77        B, T = seq_ids.shape78        positions = torch.arange(T, device=seq_ids.device).unsqueeze(0)  # [1, T]79        x = self.item_emb(seq_ids) + self.pos_emb(positions)             # [B, T, dim]8081        # 因果 mask：位置 t 只能关注 ≤ t 的位置（上三角设为 -inf）82        causal_mask = nn.Transformer.generate_square_subsequent_mask(T, device=seq_ids.device)83        out = self.transformer_layer(x, src_mask=causal_mask, is_causal=True)  # [B, T, dim]84        return out8586    def forward(self, seq_ids: torch.Tensor, pos_ids: torch.Tensor,87                neg_ids: torch.Tensor) -> torch.Tensor:88        """89        seq_ids: [B, T]   输入序列90        pos_ids: [B]      正样本（下一个物品）91        neg_ids: [B]      负样本92        返回 BCE loss93        """94        out = self._encode(seq_ids)         # [B, T, dim]95        user_repr = out[:, -1, :]           # [B, dim]  最后一个位置的表示9697        v_pos = self.item_emb(pos_ids)      # [B, dim]98        v_neg = self.item_emb(neg_ids)      # [B, dim]99100        s_pos = (user_repr * v_pos).sum(-1) # [B]101        s_neg = (user_repr * v_neg).sum(-1) # [B]102103        loss = F.binary_cross_entropy_with_logits(104            s_pos, torch.ones_like(s_pos)105        ) + F.binary_cross_entropy_with_logits(106            s_neg, torch.zeros_like(s_neg)107        )108        return loss109110    def train_model(111        self,112        user_seqs: dict[int, list[int]],113        epochs: int = 10,114        neg_k: int = 1,115        batch_size: int = 64,116    ) -> None:117        """Leave-one-out 训练：最后一个物品为正样本，前面为序列输入。"""118        samples: list[tuple[list[int], int]] = []119        for uid, seq in user_seqs.items():120            if len(seq) < 2:121                continue122            for t in range(1, len(seq)):123                samples.append((seq[:t], seq[t]))124125        optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.reg)126127        for epoch in range(epochs):128            random.shuffle(samples)129            total_loss = 0.0130            n_batches = 0131132            for start in range(0, len(samples), batch_size):133                chunk = samples[start:start + batch_size]134                # Truncate / pad context sequences to max_len135                contexts = []136                for ctx, _ in chunk:137                    ctx = ctx[-self.max_len:]138                    contexts.append(ctx)139                max_t = max(len(c) for c in contexts)140                # Left-pad with item 0 (padding token)141                padded = torch.zeros(len(chunk), max_t, dtype=torch.long)142                for i, ctx in enumerate(contexts):143                    padded[i, max_t - len(ctx):] = torch.tensor(ctx, dtype=torch.long)144145                pos_ids = torch.tensor([p for _, p in chunk], dtype=torch.long)146                neg_ids = torch.tensor(147                    [self._sample_neg(p, self.n_items) for _, p in chunk],148                    dtype=torch.long,149                )150151                optimizer.zero_grad()152                loss = self(padded, pos_ids, neg_ids)153                loss.backward()154                optimizer.step()155156                total_loss += loss.item()157                n_batches += 1158159            if epoch % 2 == 0:160                print(f"  epoch {epoch:3d}  avg_loss={total_loss / max(n_batches, 1):.4f}")161162    @staticmethod163    def _sample_neg(pos: int, n_items: int) -> int:164        neg = random.randint(0, n_items - 1)165        while neg == pos:166            neg = random.randint(0, n_items - 1)167        return neg168169    def recommend(self, seq: list[int], top_k: int = 10,170                  exclude: set[int] | None = None) -> list[tuple[int, float]]:171        """给定序列，推荐下一个物品。"""172        exclude = exclude or set(seq)173        with torch.no_grad():174            seq_t = torch.tensor(seq[-self.max_len:], dtype=torch.long).unsqueeze(0)  # [1, T]175            out = self._encode(seq_t)               # [1, T, dim]176            user_repr = out[0, -1, :]               # [dim]177            scores = self.item_emb.weight @ user_repr  # [n_items]178            for iid in exclude:179                scores[iid] = -1e9180            top_vals, top_idx = scores.topk(top_k)181        return [(int(i), float(s)) for i, s in zip(top_idx.tolist(), top_vals.tolist())]182183184# ── 快速验证 ──────────────────────────────────────────────────────────────────185if __name__ == "__main__":186    random.seed(42)187    torch.manual_seed(42)188    n_items = 15189    user_seqs = {190        u: random.sample(range(n_items), k=random.randint(4, 8))191        for u in range(8)192    }193194    model = SASRec(n_items, dim=16, max_len=10, lr=0.05)195    print("SASRec 训练：")196    model.train_model(user_seqs, epochs=6)197198    seq = user_seqs[0]199    recs = model.recommend(seq, top_k=5)200    print(f"\n用户序列 {seq} 的 Top-5 推荐：")201    for iid, s in recs:202        print(f"  item {iid}: {s:.4f}")