02 · Two towers: source code, tensors, gradients, and real training

# English companions live one directory deeper; keep the original experiment code unchanged.
from pathlib import Path as _TranslationPath
import os as _translation_os
import sys as _translation_sys
_translation_root = next((p for p in [_TranslationPath.cwd(), *_TranslationPath.cwd().parents]
                          if (p / "reclab").is_dir()), None)
if _translation_root is None:
    raise RuntimeError("Open this notebook from inside the Search Recommendation Ads Lab checkout.")
_translation_os.chdir(_translation_root)
if str(_translation_root) not in _translation_sys.path:
    _translation_sys.path.insert(0, str(_translation_root))

from pathlib import Path
import sys
import inspect
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from IPython.display import display, Code

PROJECT = next(p for p in [Path.cwd(), *Path.cwd().parents]
               if (p / "reclab/models.py").exists())
if str(PROJECT) not in sys.path:
    sys.path.insert(0, str(PROJECT))
from reclab.data import Dataset
from reclab.paths import ARTIFACTS, PROCESSED, read_json

plt.rcParams.update({"figure.figsize": (9, 3.8), "font.size": 11,
    "axes.spines.top": False, "axes.spines.right": False,
    "axes.labelcolor": "#182b2b", "text.color": "#182b2b",
    "axes.facecolor": "white", "figure.facecolor": "white",
    "grid.color": "#d4ddda", "axes.axisbelow": True})
dataset = Dataset()
print("Project:", PROJECT)
print("Dataset version:", dataset.manifest["dataset_version"])
print("Python:", sys.executable)
assert Path(sys.executable).parent == PROJECT / ".venv/bin", "请选择 Recommendation Lab 内核"

Project: /Users/lizuolin/local/recommendation-lab
Dataset version: 2772cab0aca8cb81
Python: /Users/lizuolin/local/recommendation-lab/.venv/bin/python

import torch
from torch import nn
from torch.nn import functional as F
import reclab.models as model_module
from reclab.models import TwoTower
torch.set_num_threads(4)
torch.manual_seed(42)
DIM = 32
HISTORY_LEN = 20
BATCH = 8
TEMPERATURE = 0.1

print(inspect.getfile(TwoTower))
display(Code(inspect.getsource(TwoTower), language="python"))

/Users/lizuolin/local/recommendation-lab/reclab/models.py

class TwoTower(nn.Module):
    def __init__(self, n_users, n_items, dim=32, use_history=True, dropout=0.1):
        super().__init__()
        self.use_history = use_history
        self.user_embedding = nn.Embedding(n_users, dim)
        self.item_embedding = nn.Embedding(n_items, dim, padding_idx=0)
        self.user_mlp = nn.Sequential(
            nn.Linear(dim * 2, dim * 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(dim * 2, dim),
        )
        self.item_mlp = nn.Sequential(
            nn.Linear(dim, dim * 2), nn.ReLU(), nn.Linear(dim * 2, dim)
        )

    def encode_user(self, users, history):
        mask = (history != 0).unsqueeze(-1)
        pooled = (self.item_embedding(history) * mask).sum(1) / mask.sum(1).clamp_min(1)
        if not self.use_history:
            pooled = torch.zeros_like(pooled)
        return F.normalize(
            self.user_mlp(torch.cat([self.user_embedding(users), pooled], dim=-1)),
            dim=-1,
        )

    def encode_item(self, items):
        return F.normalize(self.item_mlp(self.item_embedding(items)), dim=-1)

    def loss(self, users, history, targets, temperature=0.1, logq=None):
        # In-batch sampled softmax. Duplicate positive IDs are accidental hits, not negatives.
        logits = (
            self.encode_user(users, history) @ self.encode_item(targets).T / temperature
        )
        if logq is not None:
            logits = logits - logq[targets][None, :]
        duplicate = targets[:, None].eq(targets[None, :])
        duplicate.fill_diagonal_(False)
        logits = logits.masked_fill(duplicate, -1e9)
        return F.cross_entropy(logits, torch.arange(len(users), device=users.device))

users_np, histories_np, targets_np, times_np = dataset.examples(HISTORY_LEN)
users = torch.from_numpy(users_np[:BATCH])
histories = torch.from_numpy(histories_np[:BATCH])
targets = torch.from_numpy(targets_np[:BATCH])
model = TwoTower(dataset.n_users, dataset.n_items, dim=DIM, use_history=True, dropout=0.0)
model.eval()  # 固定本次数学诊断；eval 不等于关闭 autograd
user_vectors = model.encode_user(users, histories)
item_vectors = model.encode_item(targets)
display(pd.DataFrame({"tensor": ["users", "histories", "targets", "user_vectors", "item_vectors"],
    "shape": [list(x.shape) for x in [users, histories, targets, user_vectors, item_vectors]]}))
print("Trainable parameters:", sum(p.numel() for p in model.parameters()))
assert torch.allclose(user_vectors.norm(dim=1), torch.ones(BATCH), atol=1e-6)

Trainable parameters: 385952

cosine = user_vectors @ item_vectors.T
duplicate = targets[:, None].eq(targets[None, :])
duplicate.fill_diagonal_(False)
logits = (cosine / TEMPERATURE).masked_fill(duplicate, -1e9)
logits.retain_grad()
labels = torch.arange(BATCH)
probabilities = logits.softmax(dim=1)
loss = F.cross_entropy(logits, labels)
model.zero_grad(set_to_none=True)
loss.backward()
expected_gradient = (probabilities.detach() - F.one_hot(labels, BATCH)) / BATCH
assert torch.allclose(logits.grad, expected_gradient, atol=1e-6)
display(pd.DataFrame(probabilities.detach().numpy(), index=users.tolist(), columns=targets.tolist()))
print("Cross-entropy:", float(loss.detach()))
print("dL/dlogit for row 0:", logits.grad[0].tolist())

Cross-entropy: 2.9391064643859863
dL/dlogit for row 0: [-0.115597203373909, 0.0005943385185673833, 0.002060179365798831, 0.002936860080808401, 0.01839778758585453, 0.08384954929351807, 0.00157161767128855, 0.006186868529766798]

fig, ax = plt.subplots(figsize=(7.5, 5.5))
image = ax.imshow(probabilities.detach().numpy(), cmap="Greens", vmin=0, vmax=1)
ax.set_xticks(range(BATCH), [str(i) for i in targets.tolist()])
ax.set_yticks(range(BATCH), [str(i) for i in users.tolist()])
ax.set(xlabel="Target item index (column)", ylabel="User index (row)",
       title="Reconstructed training batch / softmax probability")
for row in range(BATCH):
    for col in range(BATCH):
        value = float(probabilities[row, col].detach())
        text = "x" if duplicate[row, col] else f"{value:.2f}"
        ax.text(col, row, text, ha="center", va="center", color="white" if value > .55 else "#182b2b")
fig.colorbar(image, ax=ax, label="Probability")
plt.tight_layout(); plt.show()

Saved notebook output · not executed on this page

<Figure size 750x550 with 2 Axes>

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
tiny_losses = []
for step in range(30):
    optimizer.zero_grad(set_to_none=True)
    tiny_loss = model.loss(users, histories, targets, TEMPERATURE)
    tiny_losses.append(float(tiny_loss.detach()))
    tiny_loss.backward()
    optimizer.step()
assert tiny_losses[-1] < tiny_losses[0]
print("Tiny-batch loss:", tiny_losses[0], "->", tiny_losses[-1])
fig, ax = plt.subplots(figsize=(8, 3))
ax.plot(range(1,31), tiny_losses, color="#006e60", marker=".")
ax.set(xlabel="Optimization step", ylabel="Cross-entropy", title="Same 8 examples / overfit sanity check", ylim=(0, max(tiny_losses)*1.1))
ax.grid(axis="y"); plt.tight_layout(); plt.show()

Saved notebook output · not executed on this page

Tiny-batch loss: 2.9391064643859863 -> 0.18255160748958588

<Figure size 800x300 with 1 Axes>

HIDDEN = 16
class MyTwoTower(TwoTower):
    def __init__(self, n_users, n_items, dim):
        super().__init__(n_users, n_items, dim, use_history=True, dropout=0.0)
        self.user_mlp = nn.Sequential(nn.Linear(2*dim, HIDDEN), nn.ReLU(), nn.Linear(HIDDEN, dim))

sandbox = MyTwoTower(dataset.n_users, dataset.n_items, DIM)
print("Edited user tower:", sandbox.user_mlp)
print("Output shape:", tuple(sandbox.encode_user(users, histories).shape))
print("Parameters:", sum(p.numel() for p in sandbox.parameters()))

Edited user tower: Sequential(
  (0): Linear(in_features=64, out_features=16, bias=True)
  (1): ReLU()
  (2): Linear(in_features=16, out_features=32, bias=True)
)
Output shape: (8, 32)
Parameters: 381296

import importlib
import reclab.train as train_module
importlib.reload(model_module)
train_module = importlib.reload(train_module)

EXPERIMENT = train_module.TrainConfig(
    name="notebook / first-full-run", dim=32, history_len=20,
    epochs=4, batch_size=256, learning_rate=0.003,
    use_history=True, logq_correction=False, seed=42,
)
new_run = train_module.train(EXPERIMENT)
print("Saved run:", new_run["run_id"])
display(pd.Series(new_run["validation"]["metrics"], name="validation score").to_frame())

20261028T023407346105 {'epoch': 1, 'loss': 5.560806380245401, 'gradient_norm_before_clip': 0.3162150839824632, 'validation_ndcg@20': 0.0007904414175060725, 'validation_recall@20': 0.000609776941282452, 'elapsed_seconds': 1.81}

20261028T023407346105 {'epoch': 2, 'loss': 5.492199623965221, 'gradient_norm_before_clip': 0.32834317195863233, 'validation_ndcg@20': 0.008092865280543902, 'validation_recall@20': 0.014427202969790203, 'elapsed_seconds': 2.38}

20261028T023407346105 {'epoch': 3, 'loss': 5.232908364200719, 'gradient_norm_before_clip': 0.5570250882509169, 'validation_ndcg@20': 0.01700333691787301, 'validation_recall@20': 0.031168373965094427, 'elapsed_seconds': 2.93}

20261028T023407346105 {'epoch': 4, 'loss': 4.913768497700725, 'gradient_norm_before_clip': 0.7106661569904273, 'validation_ndcg@20': 0.017804264366807675, 'validation_recall@20': 0.03713369853338916, 'elapsed_seconds': 3.53}

Saved run: 20261028T023407346105

curve = pd.DataFrame(new_run["curve"])
fig, axes = plt.subplots(1, 2, figsize=(10, 3.6))
axes[0].plot(curve.epoch, curve.loss, color="#006e60", marker="o")
axes[0].set(title="Training sampled-softmax", xlabel="Epoch", ylabel="Cross-entropy", ylim=(0, curve.loss.max()*1.1))
axes[1].plot(curve.epoch, curve["validation_ndcg@20"], color="#006e60", marker="o")
axes[1].set(title=f"Validation / {new_run['validation']['audit']['evaluated_users']} users", xlabel="Epoch", ylabel="NDCG@20 score", ylim=(0, max(curve["validation_ndcg@20"].max()*1.2, .001)))
for ax in axes:
    ax.set_xticks(curve.epoch)
    ax.grid(axis="y")
plt.tight_layout(); plt.show()

Saved notebook output · not executed on this page

<Figure size 1000x360 with 2 Axes>