Привет, Хабр!

Сегодня разберём RecBole — универсальный фреймворк на PyTorch, который отвечает на три насущных вопроса любого ML-инженера рекомендаций:

• Как быстро обкатать десятки алгоритмов (от классического MF до SASRec и KGAT) на собственном датасете — без сотни скриптов?

• Как хранить все настройки в одном YAML, а не в трёх сотнях аргументов CLI?

• Как получить честное сравнение метрик и сразу вынести лучший чекпоинт в прод?

Рассмотрим подробнее под катом.

Установка и подготовка данных

pip install recbole>=1.2
python -m recbole.quick_start.run_recbole --model=BPR --dataset=ml-1m

У RecBole есть встроенная заготовка датасетов (ml-1m, yelp, amazon-*). Свой датасет кидаем в папку dataset/<name>/ в формате Atomic Files:

Parquet читается быстрее, но RecBole «проглатывает» и CSV.

Автоматический сплит

В recbole.yaml достаточно:

split_ratio: [0.8, 0.1, 0.1]   # train/valid/test
group_by_user: True            # чтобы у каждого юзера были все статусы

Всё, никаких ручных датафреймов на pandas.

Разбираемся с API

run_recbole

from recbole.quick_start import run_recbole

run_recbole(
    model='LightGCN',          # любая из 90+ моделей
    dataset='ml-1m',           # или путь к своему набору
    config_dict={              # приоритет над YAML и CLI
        'epochs': 50,
        'topk': 10,
        'neg_sampling': {'uniform': 1},
        'seed': 42,            # чтобы метрики не «плавали»
    }
)

Что происходит под капотом

Хотите логировать в W&B — добавьте wandb: True в YAML. Нужен mixed-precision — train_stage: fp16. Гиперпараметры через CLI: python run_recbole.py --learning_rate=5e-4 --dropout_prob=0.3.

Гранулярный контроль

Иногда однострочник — роскошь, и нужен доступ к каждому объекту. Тогда:

from recbole.config import Config
from recbole.data import create_dataset, data_preparation
from recbole.utils import init_seed
from recbole.model.general_recommender import LightGCN
from recbole.trainer import Trainer

# 1. Конфиг из файла + CLI
config = Config(model='LightGCN', dataset='ml-1m')      # читает recbole.yaml
config['epochs'] = 30                                   # оверрайд «на лету»

# 2. Dataset
init_seed(config['seed'])
dataset = create_dataset(config)                        # <RecDataset 1 1000209>

# 3. Sampler / Dataloader
train_data, valid_data, test_data = data_preparation(config, dataset)

# 4. Модель
model = LightGCN(config, dataset).to(config['device'])

# 5. Тренер
trainer = Trainer(config, model)
best_valid_score, best_valid_result = trainer.fit(
    train_data, valid_data, saved=True, show_progress=True)

score, result = trainer.evaluate(test_data, load_best_model=True)
print(result)     # {'Recall@10': 0.1627, 'NDCG@10': 0.0894, ...}

Config

# recbole.yaml (кусочек)
MODEL_TYPE: Sequential     # автоматически подскажет, что у модели есть max_seq_length
epochs: 40
neg_sampling:
  dynamic: 1
eval_args:
  mode: full                # full-sort evaluation
  order: RO                 # рейтинг -> онлайн
  split: {'RS': [0.8,0.1,0.1]}
checkpoint_dir: ./saved/
wandb: True

Переопределение при импорте:

cfg = Config(model='SASRec', dataset='ml-1m',
             config_dict={'epochs': 10, 'dropout_prob': 0.2})

Доступ к параметрам — по ключу: cfg['topk'], cfg.final_config_dict — готовый словарь для логирования.

Dataset и друзья

from recbole.data import Dataset
dset = Dataset(config)      # наследник torch.utils.data.Dataset
len(dset.field2type)        # {'user_id': 'token', 'item_id': 'token', ...}

• Custom поля — добавьте колонку в .inter и опишите тип в YAML:

  FIELD_TYPES: {'price': float, 'brand': token}

• Sequence -> unrolled. Для последовательных моделей (SASRec, GRU4Rec) RecBole сам создаёт hist_seq и target_item.

• Lazy loading для >10 GB дат:

  lazy_loading: True

Самплеры и “кормушки”

from recbole.data import (
    create_samplers, create_dataloader, data_preparation
)

samplers = create_samplers(config, dataset)       # TrainSampler / FullSortSampler
train_loader, valid_loader, test_loader = create_dataloader(
    config, dataset, samplers)

• Популярный негатив:

  neg_sampling:
    popularity: 1

  Внутри PopularitySampler — item-frequency softmax.

• Dynamic Sampler считает свежие негативы каждую эпоху, спасая от информации-leakage.

Пишем свою модель

from recbole.model.abstract_recommender import GeneralRecommender
from recbole.model.loss import BPRLoss
import torch.nn as nn
import torch

class MyDotMF(GeneralRecommender):
    def __init__(self, config, dataset):
        super().__init__(config, dataset)
        self.embedding_size = config['embedding_size']
        self.user_embedding = nn.Embedding(
            dataset.num(self.USER_ID), self.embedding_size)
        self.item_embedding = nn.Embedding(
            dataset.num(self.ITEM_ID), self.embedding_size)
        self.loss_fct = BPRLoss()

    def forward(self, interaction):
        user = interaction[self.USER_ID]
        pos_item = interaction[self.ITEM_ID]
        user_e = self.user_embedding(user)
        item_e = self.item_embedding(pos_item)
        scores = (user_e * item_e).sum(-1)
        return scores

    def calculate_loss(self, interaction):
        pos_score = self.forward(interaction)
        neg_items = interaction[self.NEG_ITEM_ID]
        neg_e = self.item_embedding(neg_items)
        neg_score = (user_e.unsqueeze(1) * neg_e).sum(-1)
        return self.loss_fct(pos_score, neg_score)

Регистрируем:

from recbole.utils import register_model
register_model('MyDotMF', MyDotMF)

Теперь в YAML достаточно model: MyDotMF.

Минимальный кейс:

mkdir -p dataset/shop
python - <<'PY'
import pandas as pd, pyarrow.parquet as pq
df = pq.read_table('orders.parquet').to_pandas()
df[['user_id','sku','ts']].to_csv(
    'dataset/shop/shop.inter', sep='\t', index=False)
PY

recbole.yaml:

field_separator: "\t"
USER_ID_FIELD: user_id
ITEM_ID_FIELD: sku
TIME_FIELD: ts

model: SASRec
epochs: 20
learning_rate: 1e-3
neg_sampling: ~
LABEL_FIELD: click
topk: 20
metrics: ['Recall', 'NDCG', 'MRR']
device: cuda

Запуск:

from recbole.quick_start import run_recbole
run_recbole(dataset='shop')

RecBole сам сделает сплит, залиогирует Recall@20, сохранит чек-пойнт и итоговый YAML в saved/SASRec-shop-<ts>/.

Фичи

Правильный neg_sampling — бесплатный буст к NDCG

neg_sampling:
  uniform: 1
# или
  popularity: 1
# или
  dynamic: 1            # поддерживается с v1.2

dynamic может давать +5 % NDCG@10 vs uniform.

Knowledge-Graph модели

Если берёте KGAT/CFKG/TransRec, добавляйте файл графа:

knowledge_graph_file: shop.kg

Формат тривиальный: head relation tail. RecBole сам построит adjacency matrix.

GPU-OOM ловушка

Параметр train_batch_size умноженный на количество GPU → ваша фактическая матрица эмбеддингов. Когда загоняете SASRec на A100 40 GB, не забывайте, что скрытая матрица self-attention растёт квадратично от max_seq_length.

train_batch_size: 512     # ок
max_seq_length: 200       # ок
n_layers: 4               # ок

Уехали в 1024×512×6 — здравствуй, CUDA OOM.

Экспорт в прод

torch.save(model.state_dict(), 'lightgcn.pt')
# inference
model = LightGCN(config, dataset)
model.load_state_dict(torch.load('lightgcn.pt', map_location='cpu'))
model.eval()

Никаких RecBole-зависимостей в рантайме: чистый PyTorch внутри Docker.

Итоги

RecBole закрывает 80 % типовых задач ресёрча и «ML-прототипов» в одном пакете: вам остаётся только решать, какую модель кормить продакшену. Да, бывают кейсы, где нужен Sparkили multi-tower архитектура под рекламу – тогда пляшем руками. Но для большинства продуктовых рекомендателей «поднять бейзлайн» быстрее RecBole сегодня мало что умеет.

Если вы работаете с рекомендательными системами или только собираетесь внедрять их в продукт, обязательно загляните в RecBole — мощный фреймворк на PyTorch, который закрывает до 80% задач ресёрча и ML‑прототипирования «из коробки». Поддержка 90+ моделей, единый YAML для всей конфигурации, автоматическая обработка данных, гибкий negative sampling и честные метрики — всё это помогает не тратить время на рутину и быстрее выходить в прод.

Чтобы разобраться в возможностях RecBole и не тратить недели на документацию, присоединяйтесь к нашему циклу открытых уроков:

• Data Science — это проще, чем кажется! — 10 июля в 18:00

• Random Forest — мощный метод ансамблирования в ML — 16 июля в 18:00

• ML для начинающих — первые шаги с Jupyter Notebook — 23 июля в 20:00

Каждый урок — это практическое погружение: от запуска бейзлайна на своём датасете до кастомизации модели и экспорта в inference. Присоединяйтесь — и проверьте на практике, насколько RecBole может упростить вашу работу.