Skip to content

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 ​

python
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 ​

python
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 architecture
  • DeepFM: A classic baseline combining categorical-feature interactions with an MLP
  • DCN / 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 ​

python
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 ​

python
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 ​

  • SequenceFeature must use pooling="concat" because DIN / DIEN / BST need the complete sequence tensor.
  • history_features and target_features must correspond one to one and share embeddings through shared_with.
  • features contains user-profile/context features. Do not put all of target_features into it again, or the inputs will be duplicated; BST will also fail because the history and target dimensions differ.
  • For BST, embed_dim must be divisible by nhead.

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 ​

python
auc = trainer.evaluate(trainer.model, test_dl)
print(f"Test AUC: {auc:.4f}")

2. ONNX Export ​

python
trainer.export_onnx("model.onnx", dynamic_batch=True)

3. Architecture Visualization ​

python
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.

Q4: Where Can I Find More Detailed Model Guides? ​