dcn_v2.py

1"""2DCN-V2（第15章 §15.5.2）34把 DCN 的向量参数化改成矩阵：5    x_{l+1} = x_0 ⊙ (W_l x_l + b_l) + x_l6参数量从 O(Ld) 升到 O(Ld²)，表达能力更强；工业上常用低秩分解7    W_l ≈ U_l V_l^T，U_l, V_l ∈ R^{d×r}8在效果与显存之间折中。910运行：python3 ch04_feature_interaction/dcn_v2.py11"""1213from __future__ import annotations1415import pathlib16import sys1718import torch19import torch.nn as nn2021sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))2223from _train_utils import (  # noqa: E40224    N_FIELDS, MLP, BaselineMLP, count_params, embedding_table, print_comparison,25    run_experiment, stack_embeddings,26)272829class MatrixCrossLayer(nn.Module):30    """31    三种参数化方式：32        rank=None  : 满矩阵 W ∈ R^{d×d}33        rank=r>0   : 低秩 W = U V，U∈R^{d×r}, V∈R^{r×d}34        rank=0     : 退化为 DCN 的向量形式 x0 ⊙ (w ⊙ xl)，即 W = w 的对角矩阵35    rank=0 的分支单独实现，避免创建 0 维 Linear 触发 PyTorch 的空张量警告。36    """3738    def __init__(self, dim: int, rank: int | None = None):39        super().__init__()40        self.rank = rank41        if rank is None:42            self.w = nn.Linear(dim, dim, bias=True)43        elif rank == 0:44            self.vec = nn.Parameter(torch.empty(dim))45            nn.init.normal_(self.vec, std=1e-3)46            self.bias = nn.Parameter(torch.zeros(dim))47        else:48            self.u = nn.Linear(dim, rank, bias=False)49            self.v = nn.Linear(rank, dim, bias=True)5051    def forward(self, x0, xl):52        if self.rank is None:53            wx = self.w(xl)54        elif self.rank == 0:55            wx = (xl * self.vec).sum(1, keepdim=True).expand_as(xl) + self.bias56        else:57            wx = self.v(self.u(xl))58        return x0 * wx + xl596061class DCNV2(nn.Module):62    def __init__(self, dim: int = 4, n_cross: int = 2, rank: int | None = None, hidden=(32,)):63        super().__init__()64        self.emb = embedding_table(dim)65        d = dim * N_FIELDS66        self.cross = nn.ModuleList([MatrixCrossLayer(d, rank) for _ in range(n_cross)])67        self.deep = MLP(d, hidden, out_dim=None, dropout=0.2)68        self.head = nn.Linear(d + hidden[-1], 1)6970    def forward(self, x):71        x0 = stack_embeddings(self.emb, x).flatten(1)72        xl = x073        for layer in self.cross:74            xl = layer(x0, xl)75        h = self.deep(x0)76        return self.head(torch.cat([xl, h], dim=1)).squeeze(-1)777879def main() -> None:80    d = 4 * N_FIELDS81    print(f"输入维度 d = {d}（10 个 field × 4 维）\n")82    print("Cross 层参数量随参数化方式的变化（2 层）：")83    print(f"  DCN 向量版      : {2 * 2 * d:,}")84    print(f"  DCN-V2 满矩阵   : {2 * d * d + 2 * d:,}（{d}x{d} 权重 + bias）")85    print(f"  DCN-V2 低秩 r=8 : {2 * (d * 8 + 8 * d + d):,}")86    print("  → 满矩阵是向量版的 d 倍；d=80 时还能承受，d=1024 时就会 OOM。\n")8788    _, a_vec, _ = run_experiment(lambda: DCNV2(rank=0), "DCN 向量版(r=0)", steps=200)89    _, a_full, _ = run_experiment(lambda: DCNV2(rank=None), "DCN-V2-full", steps=200)90    _, a_base, _ = run_experiment(BaselineMLP, "Baseline", steps=200)91    _, a_lr, _ = run_experiment(lambda: DCNV2(rank=8), "DCN-V2-lowrank(r=8)", steps=200)9293    print_comparison([94        ("Baseline MLP", a_base, count_params(BaselineMLP())),95        ("DCN 向量版（r=0 特例）", a_vec, count_params(DCNV2(rank=0))),96        ("DCN-V2 低秩 r=8", a_lr, count_params(DCNV2(rank=8))),97        ("DCN-V2 满矩阵", a_full, count_params(DCNV2(rank=None))),98    ])99    print("\n读法：满矩阵参数量最多，但不一定显著更好——正文强调 DCN-V2 的工业价值"100          "在于'把表达能力与成本一起考虑'（§15.5.2）。")101102103if __name__ == "__main__":104    main()