Custom Models in PyTorch

Custom Models in PyTorch#

Research a new recommender in plain PyTorch and get training, evaluation, serving and ONNX export from CoreRec. You write an nn.Module; CoreRec handles id mapping, negative sampling, the training loop, devices (CUDA, Apple MPS, CPU), seen-item filtering and persistence.

The contract#

Your module implements one method:

forward(query, items) -> scores
  query   LongTensor [B]           user indices          (inputs="user")
          LongTensor [B, max_len]  item-index histories,  (inputs="history")
                                   oldest first, left-padded with 0
  items   LongTensor [B, K]        item indices
  scores  FloatTensor [B, K]       higher = better

Item indices run 1..n_items and 0 is padding, so item tables are nn.Embedding(n_items + 1, d, padding_idx=0). User indices run 0..n_users-1. CoreRec builds your module as MyModel(n_users=..., n_items=..., max_len=..., **module_kwargs), passing only the arguments your __init__ declares.

Optional: score_all(query) -> [B, n_items + 1] scores the whole catalogue in one call (column 0 is ignored). Without it, forward is called on every item.

A first model#

import numpy as np
import pandas as pd
import torch.nn as nn
from corerec.nn import Recommender


class DotModel(nn.Module):
    def __init__(self, n_users, n_items, dim=32):
        super().__init__()
        self.users = nn.Embedding(n_users, dim)
        self.items = nn.Embedding(n_items + 1, dim, padding_idx=0)
        # nn.Embedding starts at N(0, 1): dot products that large stall BPR
        nn.init.normal_(self.users.weight, std=0.05)
        nn.init.normal_(self.items.weight, std=0.05)

    def forward(self, users, items):
        return (self.users(users).unsqueeze(1) * self.items(items)).sum(-1)


# toy data with structure: user u mostly picks from item cluster u % 20
rng = np.random.default_rng(0)
users = rng.integers(0, 200, 4000)
df = pd.DataFrame({"user_id": users,
                   "item_id": (users % 20) * 15 + rng.integers(0, 15, 4000)})
test = df.groupby("user_id").tail(2)          # hold out each user's last 2
train = df.drop(test.index)

rec = Recommender(DotModel, {"dim": 32}, loss="bpr", epochs=10, lr=0.01)
rec.fit(train)
print(rec.recommend(0, top_k=10))
print(rec.history_)       # training loss per epoch

fit takes a DataFrame (user_id, item_id, optional rating and timestamp) or parallel lists. Rows with rating <= 0 are dropped, and a timestamp column sets event order.

Sequential models#

With inputs="history", every event after a user’s first becomes a training example, predicted from the events before it:

from corerec.nn import SASRecBlock, causal_mask


class TinySeq(nn.Module):
    def __init__(self, n_items, max_len, dim=32):
        super().__init__()
        self.items = nn.Embedding(n_items + 1, dim, padding_idx=0)
        self.pos = nn.Embedding(max_len, dim)
        self.block = SASRecBlock(dim, 1, 0.1)

    def forward(self, history, items):
        x = self.items(history) + self.pos.weight[: history.shape[1]]
        x = self.block(x, causal_mask(history.shape[1], history.device), None)
        return (x[:, -1].unsqueeze(1) * self.items(items)).sum(-1)


train = train.assign(timestamp=np.arange(len(train)))
seq = Recommender(TinySeq, inputs="history", loss="sampled_softmax",
                  num_negatives=50, max_len=20, epochs=5)
seq.fit(train)

Building blocks#

corerec.nn

What it is

SASRecBlock(d, heads, dropout)

Causal self-attention block (Kang & McAuley 2018)

HSTUBlock(d, heads, dqk, dv, dropout, max_len, time_buckets, use_time)

HSTU block (Zhai et al., ICML 2024)

causal_mask(n)

The [n, n] mask both blocks take

CrossLayer(dim)

DCN cross layer

FMInteraction()

Factorization-machine pairwise term over [B, F, d]

MLP(dims, dropout)

Linear/ReLU/Dropout stack

bpr_loss, bce_loss, sampled_softmax_loss

Losses over [B, 1 + K] scores, positive in column 0

These are the same blocks CoreRec’s HSTU, SASRec and DCN models use. corerec.nn.models has two complete templates, MatrixFactorization and SequentialTransformer. A custom loss is any callable(scores) -> loss over the [B, 1 + K] matrix: Recommender(MyModel, loss=my_loss).

Everything else works#

The wrapped model is an ordinary CoreRec recommender:

from corerec.evaluation import Evaluator
from corerec.serving import ModelLoader, ModelServer

truth = test.groupby("user_id").item_id.apply(list).to_dict()
print(Evaluator(metrics=["ndcg@10", "recall@10"]).evaluate(rec, truth))

rec.save("dot.pt")
same = ModelLoader().load("dot.pt")     # rebuilds DotModel from its import path
server = ModelServer(same)              # server.run() serves it over HTTP

corerec.export.to_onnx(rec, "dot.onnx") exports it for ONNX Runtime (see ONNX Export).

Saving records the module’s import path. A module defined in a notebook or script reloads fine in that same session; to load it in another process, put it in an importable file or pass it explicitly: Recommender.load("dot.pt", module_cls=DotModel).

Contents