two_tower.py

1"""2双塔召回模型（第10章 §10.3.2）34核心思想：5- 用户塔和物品塔各为一个 MLP，分别将用户特征和物品特征映射到同一语义空间6- 训练用 In-Batch 负采样（同 batch 内其他物品作为负样本）7- 线上：提前计算物品 Embedding 存入向量数据库，查询用最近邻检索89本实现简化特征为 ID Embedding + MLP，使用 PyTorch autograd。10"""1112import math13import random1415import torch16import torch.nn as nn17import torch.nn.functional as F18import torch.optim as optim192021def _build_tower(emb_dim: int, hidden_dims: list, output_dim: int) -> nn.Sequential:22    """构建单塔 MLP：emb_dim → hidden_dims → output_dim（ReLU 激活，最后一层线性输出）。"""23    layer_dims = [emb_dim] + hidden_dims + [output_dim]24    layers = []25    for i, (in_d, out_d) in enumerate(zip(layer_dims[:-1], layer_dims[1:])):26        layers.append(nn.Linear(in_d, out_d))27        if i < len(layer_dims) - 2:   # 最后一层不加激活28            layers.append(nn.ReLU())29    return nn.Sequential(*layers)303132class TwoTowerRetrieval(nn.Module):33    """34    双塔召回模型。3536    特征：用户/物品均用 ID Embedding（可扩展为其他特征）。37    训练：In-Batch 负采样 Softmax loss（温度系数 τ）。38    相似度：余弦相似度。39    """4041    def __init__(42        self,43        n_users: int,44        n_items: int,45        emb_dim: int = 32,46        hidden_dims: list = None,47        output_dim: int = 32,48        temperature: float = 0.1,49    ):50        super().__init__()51        hidden_dims = hidden_dims or [64]52        self.n_items = n_items53        self.temperature = temperature5455        self.user_emb = nn.Embedding(n_users, emb_dim)56        self.item_emb = nn.Embedding(n_items, emb_dim)5758        self.user_tower = _build_tower(emb_dim, hidden_dims, output_dim)59        self.item_tower = _build_tower(emb_dim, hidden_dims, output_dim)6061    def encode_users(self, user_ids: torch.Tensor) -> torch.Tensor:62        """返回 L2 归一化的用户向量 (B, output_dim)。"""63        x = self.user_emb(user_ids)64        out = self.user_tower(x)65        return F.normalize(out, dim=-1)6667    def encode_items(self, item_ids: torch.Tensor) -> torch.Tensor:68        """返回 L2 归一化的物品向量 (B, output_dim)。"""69        x = self.item_emb(item_ids)70        out = self.item_tower(x)71        return F.normalize(out, dim=-1)7273    def forward(self, user_ids: torch.Tensor, item_ids: torch.Tensor) -> torch.Tensor:74        """75        In-Batch Sampled Softmax：构造相似度矩阵，对角线为正样本。76        返回 cross_entropy loss。77        """78        u_emb = self.encode_users(user_ids)   # (B, d)79        i_emb = self.encode_items(item_ids)   # (B, d)80        # 相似度矩阵 (B, B)81        sim = u_emb @ i_emb.T / self.temperature82        labels = torch.arange(len(user_ids))83        return F.cross_entropy(sim, labels)8485    def train_model(86        self,87        interactions: list,88        epochs: int = 10,89        batch_size: int = 16,90        lr: float = 0.01,91    ) -> None:92        optimizer = optim.Adam(self.parameters(), lr=lr)93        for epoch in range(epochs):94            random.shuffle(interactions)95            total_loss = 0.096            n_batches = 097            self.train()98            for start in range(0, len(interactions), batch_size):99                batch = interactions[start: start + batch_size]100                if len(batch) < 2:101                    continue102                uids = torch.tensor([u for u, _ in batch], dtype=torch.long)103                iids = torch.tensor([i for _, i in batch], dtype=torch.long)104                optimizer.zero_grad()105                loss = self.forward(uids, iids)106                loss.backward()107                optimizer.step()108                total_loss += loss.item()109                n_batches += 1110            if epoch % 2 == 0:111                print(f"  epoch {epoch:3d}  avg_loss={total_loss / max(n_batches, 1):.4f}")112113    def get_item_embeddings(self) -> torch.Tensor:114        """获取所有物品的 Embedding（用于离线构建向量索引）。"""115        self.eval()116        with torch.no_grad():117            all_ids = torch.arange(self.n_items)118            return self.encode_items(all_ids)119120    def recommend(self, uid: int, top_k: int = 10, exclude: set = None) -> list:121        exclude = exclude or set()122        self.eval()123        with torch.no_grad():124            u_emb = self.encode_users(torch.tensor([uid]))   # (1, d)125            item_ids = [i for i in range(self.n_items) if i not in exclude]126            i_tensor = torch.tensor(item_ids, dtype=torch.long)127            i_embs = self.encode_items(i_tensor)             # (n, d)128            scores = (u_emb @ i_embs.T).squeeze(0).tolist()129            pairs = list(zip(item_ids, scores))130            pairs.sort(key=lambda x: x[1], reverse=True)131            return pairs[:top_k]132133134# ── 快速验证 ──────────────────────────────────────────────────────────────────135if __name__ == "__main__":136    random.seed(42)137    torch.manual_seed(42)138    n_users, n_items = 10, 20139    interactions = list(set(140        (random.randint(0, n_users - 1), random.randint(0, n_items - 1))141        for _ in range(80)142    ))143144    model = TwoTowerRetrieval(145        n_users, n_items, emb_dim=16, hidden_dims=[16], output_dim=16, temperature=0.1146    )147    print("双塔召回训练：")148    model.train_model(interactions, epochs=6, batch_size=8, lr=0.05)149150    uid = 0151    seen = {iid for u, iid in interactions if u == uid}152    recs = model.recommend(uid, top_k=5, exclude=seen)153    print(f"\n用户 {uid} Top-5 推荐：")154    for iid, s in recs:155        print(f"  item {iid}: {s:.4f}")