tdssm.py

1"""2TDSSM 时序双塔粗排模型（第11章 §11.1.2）34核心扩展（相对 DSSM）：5- 用户塔分为静态兴趣（MLP）和动态兴趣（GRU 序列编码）两部分6- 两者融合后与物品 Embedding 做余弦相似度7- 体现长期兴趣 + 短期兴趣解耦建模的思想89本实现：静态 MLP（长期兴趣）+ nn.GRU（短期兴趣），PyTorch autograd。10"""1112import math13import random14import collections15from typing import List, Tuple, Dict1617import torch18import torch.nn as nn19import torch.nn.functional as F20import torch.optim as optim212223class TDSSM(nn.Module):24    """25    时序 DSSM（TDSSM）。2627    用户塔 = 静态 MLP（长期兴趣）+ GRU（短期兴趣）→ element-wise 相加融合28    物品塔 = MLP29    融合方式：element-wise 相加（可换成 concat 或加权乘）30    """3132    def __init__(33        self,34        n_users: int,35        n_items: int,36        emb_dim: int = 32,37        static_hidden: int = 32,38        lstm_hidden: int = 32,39        output_dim: int = 32,40        temperature: float = 0.1,41    ):42        super().__init__()43        self.n_items = n_items44        self.output_dim = output_dim45        self.temperature = temperature4647        self.user_static_emb = nn.Embedding(n_users, emb_dim)48        self.item_emb = nn.Embedding(n_items, emb_dim)4950        # 静态用户塔（长期兴趣 MLP）51        self.static_mlp = nn.Sequential(52            nn.Linear(emb_dim, static_hidden),53            nn.ReLU(),54            nn.Linear(static_hidden, output_dim),55        )5657        # 短期兴趣：GRU 编码行为序列58        self.gru = nn.GRU(emb_dim, lstm_hidden, batch_first=True)59        self.gru_proj = nn.Linear(lstm_hidden, output_dim)6061        # 物品塔 MLP62        self.item_mlp = nn.Sequential(63            nn.Linear(emb_dim, output_dim),64        )6566        # 初始化 Embedding67        scale = math.sqrt(1.0 / emb_dim)68        nn.init.uniform_(self.user_static_emb.weight, -scale, scale)69        nn.init.uniform_(self.item_emb.weight, -scale, scale)7071    def encode_user(self, uid: torch.Tensor, seq_ids: torch.Tensor, seq_len: int) -> torch.Tensor:72        """73        Args:74            uid:     [B] 用户 ID（用于静态特征）75            seq_ids: [B, T] 用户近期行为序列（物品 ID），不足 T 的补 076            seq_len: 实际最大序列长度（用于 packed / mean pooling）77        Returns:78            user_emb: [B, output_dim]79        """80        # 静态兴趣81        e_static = self.static_mlp(self.user_static_emb(uid))  # [B, D]8283        # 短期兴趣（GRU）84        if seq_len > 0:85            seq_emb = self.item_emb(seq_ids)          # [B, T, emb_dim]86            _, h_n = self.gru(seq_emb)                 # h_n: [1, B, lstm_hidden]87            e_temp = self.gru_proj(h_n.squeeze(0))     # [B, output_dim]88        else:89            e_temp = torch.zeros(uid.size(0), self.output_dim)9091        # 融合（element-wise 相加）92        return e_static + e_temp9394    def encode_item(self, iid: torch.Tensor) -> torch.Tensor:95        return self.item_mlp(self.item_emb(iid))  # [B, output_dim] or [N, output_dim]9697    def forward(98        self,99        uid: torch.Tensor,100        seq_ids: torch.Tensor,101        pos_iid: torch.Tensor,102        neg_iid: torch.Tensor,103    ) -> torch.Tensor:104        """BPR loss: − log σ(s_pos − s_neg)"""105        ue = self.encode_user(uid, seq_ids, seq_ids.size(1))  # [B, D]106        pe = self.encode_item(pos_iid)                         # [B, D]107        ne = self.encode_item(neg_iid)                         # [B, D]108        s_pos = F.cosine_similarity(ue, pe, dim=1)109        s_neg = F.cosine_similarity(ue, ne, dim=1)110        loss = -F.logsigmoid(s_pos - s_neg).mean()111        return loss112113    def pre_rank(114        self, uid: int, seq: List[int], candidates: List[int], top_k: int = 20115    ) -> List[Tuple[int, float]]:116        """保留接口：pre_rank(uid, seq, candidates, top_k) -> List[Tuple[int, float]]"""117        with torch.no_grad():118            u = torch.tensor([uid])119            T = len(seq)120            if T > 0:121                seq_t = torch.tensor([seq]).long()  # [1, T]122            else:123                seq_t = torch.zeros(1, 1, dtype=torch.long)124            ue = self.encode_user(u, seq_t, T)  # [1, D]125            cand_t = torch.tensor(candidates)126            ie = self.encode_item(cand_t)         # [N, D]127            sims = F.cosine_similarity(ue, ie, dim=1)128            scores = [(candidates[j], sims[j].item()) for j in range(len(candidates))]129        scores.sort(key=lambda x: x[1], reverse=True)130        return scores[:top_k]131132133def _make_batch(134    samples: List[Tuple[int, List[int], int]],135    n_items: int,136    user_pos: Dict[int, set],137    max_seq_len: int = 20,138) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:139    uids, seqs, pos_ids, neg_ids = [], [], [], []140    for uid, context, pos_id in samples:141        neg_id = random.randint(0, n_items - 1)142        while neg_id in user_pos[uid]:143            neg_id = random.randint(0, n_items - 1)144        uids.append(uid)145        pos_ids.append(pos_id)146        neg_ids.append(neg_id)147        # pad / truncate sequence148        ctx = context[-max_seq_len:]149        pad = [0] * (max_seq_len - len(ctx))150        seqs.append(pad + ctx)151    return (152        torch.tensor(uids),153        torch.tensor(seqs, dtype=torch.long),154        torch.tensor(pos_ids),155        torch.tensor(neg_ids),156    )157158159def train(160    model: TDSSM,161    user_seqs: Dict[int, List[int]],162    epochs: int = 10,163    lr: float = 0.01,164) -> None:165    optimizer = optim.Adam(model.parameters(), lr=lr)166    user_pos: Dict[int, set] = collections.defaultdict(set)167    for uid, seq in user_seqs.items():168        for i in seq:169            user_pos[uid].add(i)170171    samples = []172    for uid, seq in user_seqs.items():173        if len(seq) < 2:174            continue175        for t in range(1, len(seq)):176            samples.append((uid, seq[:t], seq[t]))177178    for epoch in range(epochs):179        random.shuffle(samples)180        total_loss = 0.0181        uids_t, seqs_t, pos_t, neg_t = _make_batch(182            samples, model.n_items, user_pos183        )184        optimizer.zero_grad()185        loss = model(uids_t, seqs_t, pos_t, neg_t)186        loss.backward()187        optimizer.step()188        total_loss = loss.item()189        if epoch % 2 == 0:190            print(f"  epoch {epoch:3d}  avg_loss={total_loss:.4f}")191192193# ── 快速验证 ──────────────────────────────────────────────────────────────────194if __name__ == "__main__":195    random.seed(42)196    torch.manual_seed(42)197198    n_users, n_items = 8, 20199    user_seqs = {u: random.sample(range(n_items), k=random.randint(5, 10)) for u in range(n_users)}200201    model = TDSSM(n_users, n_items, emb_dim=16, static_hidden=16, lstm_hidden=16, output_dim=16)202    print("TDSSM 训练：")203    train(model, user_seqs, epochs=6, lr=0.05)204205    uid = 0206    seq = user_seqs[uid]207    candidates = [i for i in range(n_items) if i not in set(seq)]208    result = model.pre_rank(uid, seq[-5:], candidates, top_k=5)209    print(f"\n用户 {uid} 粗排 Top-5（最近 5 个行为作为短期序列）：")210    for iid, s in result:211        print(f"  item {iid}: {s:.4f}")