RecBole — «комбайн» на PyTorch для любых рекомендаций

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

Сегодня разберём 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 может упростить вашу работу.