Ranking Model Tutorial
This tutorial focuses on the shared training workflow for ranking scenarios: data preparation, feature definition, trainer usage, evaluation, and common extensions. The basic examples use the built-in Criteo sample data, while the sequence-model section uses the Amazon Electronics sample data.
I. Basic Ranking Pipeline (Criteo)
1. Data Preparation and Feature Processing
import pandas as pd
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
from torch_rechub.basic.features import DenseFeature, SparseFeature
from torch_rechub.utils.data import DataGenerator
df = pd.read_csv("examples/ranking/data/criteo/criteo_sample.csv")
# Ranking baselines generally follow this "continuous features + categorical features" input pattern.
dense_features = [f"I{i}" for i in range(1, 14)]
sparse_features = [f"C{i}" for i in range(1, 27)]
# Keep missing-value handling aligned with the repository examples to avoid reproduction differences.
df[sparse_features] = df[sparse_features].fillna("-996")
df[dense_features] = df[dense_features].fillna(0)
# Normalize continuous features and encode categorical features as discrete IDs.
scaler = MinMaxScaler()
df[dense_features] = scaler.fit_transform(df[dense_features])
for feat in sparse_features:
encoder = LabelEncoder()
df[feat] = encoder.fit_transform(df[feat].astype(str))
# These Feature objects describe how each column should be fed into the model.
dense_feas = [DenseFeature(name) for name in dense_features]
sparse_feas = [SparseFeature(name, vocab_size=df[name].nunique(), embed_dim=16) for name in sparse_features]
x = df.drop(columns=["label"])
y = df["label"]
# DataGenerator automatically splits the training, validation, and test sets according to split_ratio.
dg = DataGenerator(x, y)
train_dl, val_dl, test_dl = dg.generate_dataloader(split_ratio=[0.7, 0.1], batch_size=256)2. Shared Training Pattern for WideDeep / DeepFM / DCN
import os
from torch_rechub.models.ranking import WideDeep, DeepFM, DCN
from torch_rechub.trainers import CTRTrainer
# Choose any one of the models.
# DeepFM is a good first ranking pipeline; the commented WideDeep / DCN variants only switch the model and do not change the data flow.
model = DeepFM(
deep_features=dense_feas + sparse_feas,
fm_features=sparse_feas,
mlp_params={"dims": [256, 128], "dropout": 0.2, "activation": "relu"},
)
# model = WideDeep(
# wide_features=sparse_feas,
# deep_features=sparse_feas + dense_feas,
# mlp_params={"dims": [256, 128], "dropout": 0.2, "activation": "relu"},
# )
# model = DCN(
# features=dense_feas + sparse_feas,
# n_cross_layers=3,
# mlp_params={"dims": [256, 128]},
# )
trainer = CTRTrainer(
model,
optimizer_params={"lr": 1e-3, "weight_decay": 1e-5},
n_epoch=2,
device="cpu", # Change to "cuda:0" for GPU.
model_path="./saved/ctr_basic",
)
# Create the save directory before training so saving the best weights after fit does not fail.
os.makedirs("./saved/ctr_basic", exist_ok=True)
trainer.fit(train_dl, val_dl)
# Passing trainer.model to evaluate evaluates the best model currently held by the trainer.
auc = trainer.evaluate(trainer.model, test_dl)
print(f"Test AUC: {auc:.4f}")3. Which Models Does This Section Apply To?
WideDeep: Quickly validate a wide + deep architectureDeepFM: A classic baseline combining categorical-feature interactions with an MLPDCN / DCNv2: Explicit feature crossing
These models all use the same DenseFeature + SparseFeature + DataGenerator + CTRTrainer training pattern.
II. Sequence Ranking Pipeline (DIN / BST)
The main difference between sequence models and basic ranking models is that you must generate historical behavior sequences and keep history_features strictly aligned with target_features.
1. Build Sequences from the Amazon Electronics Sample Data
import pandas as pd
from torch_rechub.basic.features import SparseFeature, SequenceFeature
from torch_rechub.utils.data import DataGenerator, df_to_dict, generate_seq_feature
data = pd.read_csv("examples/ranking/data/amazon-electronics/amazon_electronics_sample.csv")
# generate_seq_feature sorts by time and generates historical item / category sequences for each sample.
train, val, test = generate_seq_feature(
data=data,
user_col="user_id",
item_col="item_id",
time_col="time",
item_attribute_cols=["cate_id"],
)
n_users = data["user_id"].max()
n_items = data["item_id"].max()
n_cates = data["cate_id"].max()
# features contains only user-profile/context features; target_features and history_features correspond one to one.
features = [
SparseFeature("user_id", vocab_size=n_users + 1, embed_dim=8),
]
target_features = [
SparseFeature("target_item_id", vocab_size=n_items + 1, embed_dim=8),
SparseFeature("target_cate_id", vocab_size=n_cates + 1, embed_dim=8),
]
history_features = [
# Sequence models must retain the complete sequence tensor here, so use concat rather than mean / sum.
SequenceFeature("hist_item_id", vocab_size=n_items + 1, embed_dim=8, pooling="concat", shared_with="target_item_id"),
SequenceFeature("hist_cate_id", vocab_size=n_cates + 1, embed_dim=8, pooling="concat", shared_with="target_cate_id"),
]
# Convert each DataFrame to a dict before passing it to DataGenerator, matching the examples/ranking workflow.
train_dict, val_dict, test_dict = df_to_dict(train), df_to_dict(val), df_to_dict(test)
train_y = train_dict.pop("label")
val_y = val_dict.pop("label")
test_y = test_dict.pop("label")
dg = DataGenerator(train_dict, train_y)
train_dl, val_dl, test_dl = dg.generate_dataloader(
x_val=val_dict,
y_val=val_y,
x_test=test_dict,
y_test=test_y,
batch_size=4096,
)2. Creating DIN / BST
import os
from torch_rechub.models.ranking import DIN, BST
from torch_rechub.trainers import CTRTrainer
model = DIN(
features=features,
history_features=history_features,
target_features=target_features,
mlp_params={"dims": [256, 128]},
attention_mlp_params={"dims": [256, 128]},
)
# model = BST(
# features=features,
# history_features=history_features,
# target_features=target_features,
# mlp_params={"dims": [256, 128]},
# nhead=8,
# dropout=0.2,
# num_layers=1,
# )
trainer = CTRTrainer(
model,
optimizer_params={"lr": 1e-3, "weight_decay": 1e-3},
n_epoch=2,
earlystop_patience=2,
device="cpu", # Change to "cuda:0" for GPU.
model_path="./saved/ctr_sequence",
)
# Again, create the save directory first so saving weights after training does not fail.
os.makedirs("./saved/ctr_sequence", exist_ok=True)
trainer.fit(train_dl, val_dl)
auc = trainer.evaluate(trainer.model, test_dl)
print(f"Test AUC: {auc:.4f}")3. Key Constraints for Sequence Models
SequenceFeaturemust usepooling="concat"because DIN / DIEN / BST need the complete sequence tensor.history_featuresandtarget_featuresmust correspond one to one and share embeddings throughshared_with.featurescontains user-profile/context features. Do not put all oftarget_featuresinto it again, or the inputs will be duplicated; BST will also fail because the history and target dimensions differ.- For
BST,embed_dimmust be divisible bynhead.
DIEN also requires a per-timestep negative-sample sequence in neg_history_features, and CTRTrainer must set loss_mode=False; you cannot simply replace the model class above. Refer directly to the DIEN Tutorial.
III. Evaluation, Export, and Visualization
1. Evaluation
auc = trainer.evaluate(trainer.model, test_dl)
print(f"Test AUC: {auc:.4f}")2. ONNX Export
trainer.export_onnx("model.onnx", dynamic_batch=True)3. Architecture Visualization
from torch_rechub.utils.visualization import visualize_model
visualize_model(model, save_path="model_architecture.png", dpi=300)Visualization requires an additional dependency:
pip install "torch-rechub[visualization]"
IV. FAQ
Q1: Why Doesn't This Page Include Complete Code for Every Ranking Model?
Ranking models can be divided into two groups:
- Basic ranking:
WideDeep / DeepFM / DCN - Sequence ranking:
DIN / DIEN / BST
Their data-preparation workflows differ. This page preserves the shared workflow, while each model page covers model-specific parameters and tuning guidance.
Q2: How Do I Switch to a GPU?
Change device="cpu" to device="cuda:0".
Q3: Why Does the Example Call os.makedirs First?
The current CTRTrainer saves weights directly to model_path and does not create the directory automatically. Create the save directory before training so the example can run as written.
