推荐系统（三）—— 深度学习基础模型

1
2
3
4
5

import torch
import torch.nn as nn

user_embedding = nn.Embedding(num_embeddings=10_000_000, embedding_dim=64)
user_vec = user_embedding(torch.LongTensor([42]))   # shape: [1, 64]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

import torch
import torch.nn as nn

class IdEmbedding(nn.Module):
    """带 Xavier 初始化和可选 padding 的 Embedding 层。"""

    def __init__(self, vocab_size: int, dim: int, padding_idx: int | None = None):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, dim, padding_idx=padding_idx)
        nn.init.xavier_uniform_(self.emb.weight)
        if padding_idx is not None:
            with torch.no_grad():
                self.emb.weight[padding_idx].zero_()

    def forward(self, ids: torch.LongTensor) -> torch.Tensor:
        return self.emb(ids)

1
2
3
4
5
6
7
8

class MultiFieldEmbedding(nn.Module):
    def __init__(self, field_dims: list[int], dim: int):
        super().__init__()
        self.tables = nn.ModuleList([IdEmbedding(v, dim) for v in field_dims])

    def forward(self, x: torch.LongTensor) -> torch.Tensor:
        # x: [batch, num_fields] -> [batch, num_fields, dim]
        return torch.stack([t(x[:, i]) for i, t in enumerate(self.tables)], dim=1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

import torch
import torch.nn as nn

class NeuMF(nn.Module):
    """NeuMF：融合 GMF 与 MLP，两条路径独立 Embedding（He 等，2017）。"""

    def __init__(
        self,
        num_users: int,
        num_items: int,
        embedding_dim: int = 32,
        mlp_layers: tuple[int, ...] = (128, 64, 32),
        dropout: float = 0.0,
    ):
        super().__init__()
        # 两条路径独立 Embedding——NeuMF 的关键设计。
        self.user_gmf = nn.Embedding(num_users, embedding_dim)
        self.item_gmf = nn.Embedding(num_items, embedding_dim)
        self.user_mlp = nn.Embedding(num_users, embedding_dim)
        self.item_mlp = nn.Embedding(num_items, embedding_dim)

        # MLP 塔，输入是拼接后的用户/物品向量。
        layers, in_dim = [], embedding_dim * 2
        for out_dim in mlp_layers:
            layers += [nn.Linear(in_dim, out_dim), nn.ReLU()]
            if dropout > 0:
                layers.append(nn.Dropout(dropout))
            in_dim = out_dim
        self.mlp = nn.Sequential(*layers)

        # 输出头看到的是 [GMF 向量；MLP 向量]。
        self.head = nn.Linear(embedding_dim + mlp_layers[-1], 1)

        for emb in (self.user_gmf, self.item_gmf, self.user_mlp, self.item_mlp):
            nn.init.normal_(emb.weight, std=0.01)

    def forward(self, users: torch.LongTensor, items: torch.LongTensor) -> torch.Tensor:
        gmf = self.user_gmf(users) * self.item_gmf(items)           # [B, d]
        mlp_in = torch.cat([self.user_mlp(users), self.item_mlp(items)], dim=-1)
        mlp_out = self.mlp(mlp_in)                                  # [B, mlp_layers[-1]]
        logits = self.head(torch.cat([gmf, mlp_out], dim=-1)).squeeze(-1)
        return torch.sigmoid(logits)


# 冒烟测试
model = NeuMF(num_users=10_000, num_items=5_000, embedding_dim=32, dropout=0.2)
users = torch.randint(0, 10_000, (8,))
items = torch.randint(0, 5_000, (8,))
labels = torch.randint(0, 2, (8,)).float()
preds = model(users, items)
loss = nn.BCELoss()(preds, labels)
print(f"preds={preds.detach().numpy().round(3)}  loss={loss.item():.4f}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

import torch
import torch.nn as nn
import torch.nn.functional as F

class YouTubeCandidateTower(nn.Module):
    """召回塔：产出可以直接进 ANN 的用户向量。"""

    def __init__(
        self,
        num_videos: int,
        num_searches: int,
        num_geos: int,
        embedding_dim: int = 64,
        user_dim: int = 256,
    ):
        super().__init__()
        self.video_emb = nn.Embedding(num_videos, embedding_dim, padding_idx=0)
        self.search_emb = nn.Embedding(num_searches, embedding_dim, padding_idx=0)
        self.geo_emb = nn.Embedding(num_geos, embedding_dim)

        in_dim = 3 * embedding_dim + 2  # 视频池化 + 搜索池化 + 地域 + 年龄 + 性别
        self.tower = nn.Sequential(
            nn.Linear(in_dim, 1024), nn.ReLU(),
            nn.Linear(1024, 512), nn.ReLU(),
            nn.Linear(512, user_dim), nn.ReLU(),
        )

    def forward(self, watched, searched, geo, age, gender):
        # watched, searched: [B, L]，padding 过；对非 padding 的位置做 mean pool。
        v = self._masked_mean(self.video_emb(watched), watched != 0)
        s = self._masked_mean(self.search_emb(searched), searched != 0)
        g = self.geo_emb(geo)
        x = torch.cat([v, s, g, age.unsqueeze(-1), gender.unsqueeze(-1)], dim=-1)
        return F.normalize(self.tower(x), dim=-1)  # 归一化的用户向量，配 cosine 检索

    @staticmethod
    def _masked_mean(emb: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        m = mask.unsqueeze(-1).float()
        return (emb * m).sum(1) / m.sum(1).clamp(min=1.0)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

import torch
import torch.nn as nn

class WideAndDeep(nn.Module):
    """Wide & Deep（Cheng 等，2016），二分类 CTR 预测。"""

    def __init__(
        self,
        wide_dim: int,                  # 稀疏 + 交叉特征向量的维度
        field_dims: list[int],          # Deep 端每个类别字段的词表大小
        embedding_dim: int = 32,
        deep_layers: tuple[int, ...] = (256, 128, 64),
        dropout: float = 0.0,
    ):
        super().__init__()
        # Wide 端：一个大的稀疏线性层。
        self.wide = nn.Linear(wide_dim, 1)

        # Deep 端：每字段 Embedding + MLP。
        self.embeddings = nn.ModuleList([nn.Embedding(v, embedding_dim) for v in field_dims])
        in_dim = len(field_dims) * embedding_dim
        layers = []
        for out_dim in deep_layers:
            layers += [nn.Linear(in_dim, out_dim), nn.ReLU()]
            if dropout > 0:
                layers.append(nn.Dropout(dropout))
            in_dim = out_dim
        layers.append(nn.Linear(in_dim, 1))
        self.deep = nn.Sequential(*layers)

    def forward(self, x_wide: torch.Tensor, x_deep: torch.LongTensor) -> torch.Tensor:
        wide_score = self.wide(x_wide).squeeze(-1)
        deep_in = torch.cat([emb(x_deep[:, i]) for i, emb in enumerate(self.embeddings)], dim=-1)
        deep_score = self.deep(deep_in).squeeze(-1)
        return torch.sigmoid(wide_score + deep_score)


# 双优化器训练——很多教程会漏掉这个生产级细节。
model = WideAndDeep(wide_dim=10_000, field_dims=[50_000, 10_000, 100, 20])
opt_wide = torch.optim.Adagrad(model.wide.parameters(), lr=0.05)
opt_deep = torch.optim.Adam(
    [p for n, p in model.named_parameters() if not n.startswith("wide.")],
    lr=1e-3, weight_decay=1e-5,
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

import numpy as np

def sample_negatives(user_history: set[int], catalog_size: int, k: int = 4) -> list[int]:
    """均匀随机采样负样本，命中用户历史则拒绝。"""
    out = []
    while len(out) < k:
        cand = int(np.random.randint(catalog_size))
        if cand not in user_history:
            out.append(cand)
    return out

1
2
3
4

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode="min", factor=0.5, patience=3
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

def train(model, train_loader, val_loader, *, epochs=50, patience=8, device="cuda"):
    model = model.to(device)
    bce = nn.BCELoss()
    opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode="min", patience=3, factor=0.5)

    best_val, bad = float("inf"), 0
    for epoch in range(1, epochs + 1):
        model.train()
        train_loss = 0.0
        for users, items, labels in train_loader:
            users, items, labels = users.to(device), items.to(device), labels.to(device)
            opt.zero_grad()
            preds = model(users, items)
            loss = bce(preds, labels)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
            train_loss += loss.item() * len(labels)
        train_loss /= len(train_loader.dataset)

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for users, items, labels in val_loader:
                users, items, labels = users.to(device), items.to(device), labels.to(device)
                val_loss += bce(model(users, items), labels).item() * len(labels)
        val_loss /= len(val_loader.dataset)
        sched.step(val_loss)

        if val_loss < best_val - 1e-4:
            best_val, bad = val_loss, 0
            torch.save(model.state_dict(), "best.pt")
        else:
            bad += 1
            if bad >= patience:
                print(f"early stop at epoch {epoch}; best val_loss={best_val:.4f}")
                break
        print(f"epoch {epoch:>3}  train={train_loss:.4f}  val={val_loss:.4f}")