dssm.py

1"""2DSSM 双塔粗排模型（第11章 §11.1.1）34粗排场景下的 DSSM：5- 用户塔和物品塔各为一个 MLP6- 训练时从精排日志中采样正/负样本7- 离线预计算物品 Embedding，线上仅计算用户 Embedding + 内积排序89本实现：PyTorch autograd，演示粗排双塔的核心结构与训练流程。10"""1112import math13import random14import collections15from typing import List, Tuple, Optional1617import torch18import torch.nn as nn19import torch.nn.functional as F20import torch.optim as optim212223class TowerMLP(nn.Module):24    """双塔共用的 MLP 结构：多层线性 + ReLU，最后一层无激活。"""2526    def __init__(self, dims: List[int]):27        super().__init__()28        layers = []29        for in_d, out_d in zip(dims[:-1], dims[1:]):30            layers.append(nn.Linear(in_d, out_d))31        self.layers = nn.ModuleList(layers)3233    def forward(self, x: torch.Tensor) -> torch.Tensor:34        h = x35        for i, layer in enumerate(self.layers):36            h = layer(h)37            if i < len(self.layers) - 1:38                h = F.relu(h)39        return h404142class DSSM(nn.Module):43    """44    粗排 DSSM 双塔模型。4546    训练策略：Sampled Softmax（In-Batch 负采样），温度系数 τ47    匹配函数：余弦相似度48    """4950    def __init__(51        self,52        n_users: int,53        n_items: int,54        emb_dim: int = 32,55        hidden_dims: Optional[List[int]] = None,56        output_dim: int = 32,57        temperature: float = 0.05,58    ):59        super().__init__()60        hidden_dims = hidden_dims or [64, 32]61        self.temperature = temperature6263        self.user_emb = nn.Embedding(n_users, emb_dim)64        self.item_emb = nn.Embedding(n_items, emb_dim)6566        dims = [emb_dim] + hidden_dims + [output_dim]67        self.user_tower = TowerMLP(dims)68        self.item_tower = TowerMLP(dims)6970        # 初始化 Embedding71        nn.init.uniform_(self.user_emb.weight, -math.sqrt(1.0 / emb_dim), math.sqrt(1.0 / emb_dim))72        nn.init.uniform_(self.item_emb.weight, -math.sqrt(1.0 / emb_dim), math.sqrt(1.0 / emb_dim))7374    def encode_user(self, uid: torch.Tensor) -> torch.Tensor:75        return self.user_tower(self.user_emb(uid))7677    def encode_item(self, iid: torch.Tensor) -> torch.Tensor:78        return self.item_tower(self.item_emb(iid))7980    def forward(self, uid: torch.Tensor, iid: torch.Tensor) -> torch.Tensor:81        """82        In-Batch Sampled Softmax：83          sim[i, j] = cosine(u_emb[i], i_emb[j]) / temperature84          loss = cross_entropy(sim, arange(B))  — 对角线为正样本85        """86        u_emb = self.encode_user(uid)   # [B, D]87        i_emb = self.encode_item(iid)   # [B, D]88        # sim[i, j] = cosine(u_emb[i], i_emb[j])89        sim = F.cosine_similarity(u_emb.unsqueeze(1), i_emb.unsqueeze(0), dim=2) / self.temperature  # [B, B]90        labels = torch.arange(len(uid))91        loss = F.cross_entropy(sim, labels)92        return loss9394    def score(self, uid: int, iid: int) -> float:95        with torch.no_grad():96            u = torch.tensor([uid])97            i = torch.tensor([iid])98            ue = self.encode_user(u)99            ie = self.encode_item(i)100            return F.cosine_similarity(ue, ie).item()101102    def pre_rank(self, uid: int, candidates: List[int], top_k: int = 20) -> List[Tuple[int, float]]:103        """粗排：对候选集打分并返回 Top-K。"""104        with torch.no_grad():105            u = torch.tensor([uid])106            ue = self.encode_user(u)  # [1, D]107            cand_t = torch.tensor(candidates)108            ie = self.encode_item(cand_t)  # [N, D]109            sims = F.cosine_similarity(ue, ie, dim=1)  # [N]110            scores = [(candidates[j], sims[j].item()) for j in range(len(candidates))]111        scores.sort(key=lambda x: x[1], reverse=True)112        return scores[:top_k]113114115def train(116    model: DSSM,117    interactions: List[Tuple[int, int]],118    epochs: int = 10,119    batch_size: int = 32,120    lr: float = 0.05,121) -> None:122    optimizer = optim.Adam(model.parameters(), lr=lr)123    for epoch in range(epochs):124        random.shuffle(interactions)125        total_loss, n_batches = 0.0, 0126        for start in range(0, len(interactions), batch_size):127            batch = interactions[start: start + batch_size]128            if len(batch) < 2:129                continue130            uids = torch.tensor([u for u, _ in batch])131            iids = torch.tensor([i for _, i in batch])132            optimizer.zero_grad()133            loss = model(uids, iids)134            loss.backward()135            optimizer.step()136            total_loss += loss.item()137            n_batches += 1138        if epoch % 2 == 0:139            print(f"  epoch {epoch:3d}  avg_loss={total_loss / max(n_batches, 1):.4f}")140141142# ── 快速验证 ──────────────────────────────────────────────────────────────────143if __name__ == "__main__":144    random.seed(42)145    torch.manual_seed(42)146147    n_users, n_items = 12, 30148    interactions = list(set(149        (random.randint(0, n_users - 1), random.randint(0, n_items - 1))150        for _ in range(100)151    ))152153    model = DSSM(n_users, n_items, emb_dim=16, hidden_dims=[32, 16], output_dim=16, temperature=0.1)154    print("DSSM 粗排模型训练：")155    train(model, interactions, epochs=6, batch_size=16, lr=0.05)156157    uid = 0158    seen = {iid for u, iid in interactions if u == uid}159    candidates = [i for i in range(n_items) if i not in seen]160    result = model.pre_rank(uid, candidates, top_k=5)161    print(f"\n用户 {uid} 粗排 Top-5：")162    for iid, s in result:163        print(f"  item {iid}: {s:.4f}")