# datawhalechina/torch-rechub architecture

> Architecture diagram of the GitHub repository datawhalechina/torch-rechub, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-25). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/datawhalechina/torch-rechub
- Repository: https://github.com/datawhalechina/torch-rechub

## Overview

1. **Repository role**
   - This is a **PyTorch-based recommender-system framework** centered on a reusable Python package, `torch_rechub`.
   - The repo is organized as a library first, with supporting **examples, tutorials, benchmarks, docs, and tests** around it.

2. **Core library boundary: `torch_rechub/`**
   - The main architecture lives under `torch_rechub`, which is split into:
     - `basic/`: shared neural-network building blocks, losses, metrics, feature helpers, callbacks, tracking, and initialization utilities.
     - `models/`: concrete recommender implementations.
     - `trainers/`: task-specific training/evaluation orchestration.
     - `utils/`: cross-cutting helpers for data, model conversion, ONNX export, visualization, matching, multitask logic, and HSTU support.
     - `serving/`: retrieval/indexing adapters for deployment-oriented lookup backends.

3. **Model subsystem**
   - Models are grouped by recommendation task:
     - **Ranking / CTR** models: DeepFM, Wide&Deep, DCN, DIN, DIEN, BST, AFM, AutoInt, FiBiNET, etc.
     - **Matching / retrieval** models: DSSM, YoutubeDNN, MIND, GRU4Rec, SASRec, NARM, STAMP, ComiRec, SINE, etc.
     - **Multi-task** models: ESMM, MMoE, PLE, AITM, SharedBottom.
     - **Generative recommendation** models: HLLM, HSTU, TIGER, RQVAE.
   - Architecturally, these are separate model families but reuse the same feature/layer/loss primitives from `basic/`.

4. **Training pipeline boundary**
   - Training is abstracted by dedicated trainers:
     - `CTRTrainer` for ranking/CTR tasks
     - `MatchTrainer` for two-tower / retrieval tasks
     - `MTLTrainer` for multi-task setups
     - `SeqTrainer` and `RQVAETrainer` for sequence/generative variants
   - The expected flow is: **dataset → dataloader → model → trainer.fit() → evaluation/export**.
   - This separation keeps task-specific optimization logic out of model definitions.

5. **Data handling subsystem**
   - Data preparation is handled through `torch_rechub/data` and `torch_rechub/utils/data.py`, with support for standardized feature conversion and dataset abstractions.
   - The repo supports both classic tabular interaction data and sequence-based data for session/retrieval models.
   - The README and examples indicate compatibility with common recommender datasets such as MovieLens, Criteo, Avazu, Census, Amazon variants, Ali-CCP, and session datasets.

6. **Serving and deployment path**
   - The repo supports deployment-oriented workflows through `torch_rechub/serving` plus ONNX utilities.
   - `utils/onnx_export.py` and `utils/quantization.py` handle model export and compression workflows.
   - `serving/annoy.py`, `serving/faiss.py`, and `serving/milvus.py` show that retrieval models can be backed by **approximate nearest neighbor / vector-store backends** for online matching.
   - The architecture therefore spans from offline training to vector-index-based serving for retrieval use cases.

7. **Experiment tooling and observability**
   - `basic/tracking.py` and the README show built-in integration points for **WandB, SwanLab, and TensorBoardX**.
   - `basic/callback.py` and visualization utilities suggest experiment hooks for monitoring, logging, and model-graph inspection.
   - This is a framework-level concern rather than embedded in individual models or trainers.

8. **Benchmarking subsystem**
   - The `benchmarks/` directory is a separate, config-driven evaluation harness.
   - It contains:
     - dataset/model/metric schemas
     - task suites for matching, ranking, and multitask benchmarks
     - baseline definitions in YAML
     - runners and reports
   - This subsystem provides reproducible comparisons across model families and is distinct from the end-user example scripts.

9. **Examples as executable documentation**
   - `examples/` contains runnable scripts organized by task:
     - `examples/ranking`
     - `examples/matching`
     - `examples/generative`
     - `examples/serving`
   - These scripts include dataset preprocessing, training entry points, and deployment demos.
   - They are the main “integration layer” showing how the package is intended to be used end to end.

10. **Docs and tutorials**
   - `docs/` is a full VitePress documentation site in English and Chinese, covering core concepts, model families, tools, serving, and benchmarks.
   - `tutorials/` contains Jupyter notebooks for step-by-step workflows such as CTR, ranking, matching, multitask learning, tracking, and serving.
   - Documentation is a first-class part of the repository, not an afterthought.

11. **Testing and quality gates**
   - `tests/` includes end-to-end tests for ranking, matching, multitask, serving, ONNX export/quantization, preprocessing, and dataset behavior.
   - The repo also includes configuration for linting, formatting, pre-commit hooks, and CI workflows under `.github/` and `config/`.
   - This indicates the project treats model correctness, exportability, and serving compatibility as part of its contract.

12. **Primary data flow for architecture diagrams**
   - A typical system flow is:
     **raw dataset or preprocessed sample → data conversion/dataloader → model family → trainer → evaluation/tracking → optional ONNX export → optional ANN/vector serving backend**
   - For benchmarks, the flow is similar but controlled by YAML configs and suite runners.
   - For examples/tutorials, the same pipeline is wrapped in task-specific scripts and notebooks.

13. **External/runtime dependencies that matter architecturally**
   - The framework is built on **Python 3.9+ and PyTorch 1.10+**.
   - Optional integrations materially affect system boundaries:
     - **ONNX** for export and runtime inference
     - **Faiss / Annoy / Milvus** for retrieval indexing and serving
     - **PyArrow** for big-data / Parquet workflows
     - **TorchView / Graphviz** for visualization
   - These are optional extension points rather than hard dependencies of the core library.

14. **Overall architectural shape**
   - The repo is best viewed as a **modular recommender-system platform** with three main layers:
     - a reusable **core library** for models, training, and utilities
     - a set of **task-specific execution paths** in examples and benchmarks
     - a **deployment/serving path** via ONNX and vector indices
   - Supporting layers for docs, tutorials, CI, and tests surround the core to keep the framework reproducible and usable.

Main entry points: [`torch_rechub/trainers`](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/trainers) (Trainers), [`torch_rechub/models`](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models) (Models), [`torch_rechub`](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub) (torch_rechub), [`examples`](https://github.com/datawhalechina/torch-rechub/tree/main/examples) (Examples)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_core["Core library"]
  node_torch_rechub["torch_rechub<br/>Python package"]
  node_basic["Basic<br/>NN primitives"]
  node_models["Models<br/>Model families"]
  node_trainers["Trainers<br/>Training orchestration"]
  node_data["Data"]
  node_utils["Utils<br/>Cross-cutting tools"]
  node_serving["Serving"]
  node_ranking_models["Ranking<br/>CTR models"]
  node_matching_models["Matching<br/>Retrieval models"]
  node_mtl_models["Multi-task<br/>MTL models"]
  node_generative_models["Generative<br/>Seq/gen models"]
  node_ctr_trainer["CTR trainer<br/>Ranking trainer<br/>[ctr_trainer.py]"]
  node_match_trainer["Match trainer<br/>Retrieval trainer<br/>[match_trainer.py]"]
  node_mtl_trainer["MTL trainer<br/>Multi-task trainer<br/>[mtl_trainer.py]"]
  node_seq_trainer["Seq trainer<br/>Sequence trainer<br/>[seq_trainer.py]"]
  node_rqvae_trainer["RQVAE trainer<br/>Generative trainer<br/>[rqvae_trainer.py]"]
end

subgraph group_workflows["Workflows"]
  node_benchmarks["Benchmarks<br/>Benchmark harness"]
  node_examples["Examples<br/>Runnable demos"]
  node_docs["Docs<br/>VitePress site"]
  node_tutorials["Tutorials<br/>Notebooks"]
end

subgraph group_quality["Quality"]
  node_tests["Tests<br/>E2E coverage"]
end

node_torch_rechub -->|"composes"| node_basic
node_torch_rechub -->|"contains"| node_models
node_torch_rechub -->|"contains"| node_trainers
node_torch_rechub -->|"contains"| node_data
node_torch_rechub -->|"contains"| node_utils
node_torch_rechub -->|"contains"| node_serving
node_models -->|"specializes"| node_ranking_models
node_models -->|"specializes"| node_matching_models
node_models -->|"specializes"| node_mtl_models
node_models -->|"specializes"| node_generative_models
node_ranking_models -->|"reuses"| node_basic
node_matching_models -->|"reuses"| node_basic
node_mtl_models -->|"reuses"| node_basic
node_generative_models -->|"reuses"| node_basic
node_trainers -->|"includes"| node_ctr_trainer
node_trainers -->|"includes"| node_match_trainer
node_trainers -->|"includes"| node_mtl_trainer
node_trainers -->|"includes"| node_seq_trainer
node_trainers -->|"includes"| node_rqvae_trainer
node_data -->|"feeds"| node_trainers
node_examples -->|"demonstrates"| node_models
node_examples -->|"drives"| node_trainers
node_examples -->|"shows"| node_serving
node_benchmarks -->|"evaluates"| node_models
node_benchmarks -->|"uses"| node_data
node_docs -->|"documents"| node_examples
node_tutorials -->|"teaches"| node_examples
node_tests -->|"verifies"| node_trainers
node_tests -->|"verifies"| node_serving
node_tests -->|"verifies"| node_utils

click node_torch_rechub "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub"
click node_basic "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/basic"
click node_models "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models"
click node_trainers "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/trainers"
click node_data "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/data"
click node_utils "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/utils"
click node_serving "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/serving"
click node_ranking_models "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/ranking"
click node_matching_models "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/matching"
click node_mtl_models "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/multi_task"
click node_generative_models "https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/generative"
click node_ctr_trainer "https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/ctr_trainer.py"
click node_match_trainer "https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/match_trainer.py"
click node_mtl_trainer "https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/mtl_trainer.py"
click node_seq_trainer "https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/seq_trainer.py"
click node_rqvae_trainer "https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/rqvae_trainer.py"
click node_benchmarks "https://github.com/datawhalechina/torch-rechub/tree/main/benchmarks"
click node_examples "https://github.com/datawhalechina/torch-rechub/tree/main/examples"
click node_docs "https://github.com/datawhalechina/torch-rechub/tree/main/docs"
click node_tutorials "https://github.com/datawhalechina/torch-rechub/tree/main/tutorials"
click node_tests "https://github.com/datawhalechina/torch-rechub/tree/main/tests"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_torch_rechub,node_basic,node_models,node_trainers,node_data,node_utils,node_serving,node_ranking_models,node_matching_models,node_mtl_models,node_generative_models,node_ctr_trainer,node_match_trainer,node_mtl_trainer,node_seq_trainer,node_rqvae_trainer toneBlue
class node_benchmarks,node_examples,node_docs,node_tutorials toneAmber
class node_tests toneMint
```

## Components

### Core library

Reusable recommender-system package and its main runtime subsystems.

- [torch_rechub](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub): Python package · Main reusable recommender framework. · `torch_rechub`
- [Basic](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/basic): NN primitives · Shared layers, activations, losses, metrics, features, callbacks, tracking, and initialization. · `torch_rechub/basic`
- [Models](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models): Model families · Ranking, matching, multi-task, and generative recommender implementations. · `torch_rechub/models`
- [Trainers](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/trainers): Training orchestration · Task-specific fit/eval loops for ranking, matching, multitask, sequence, and generative training. · `torch_rechub/trainers`
- [Data](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/data): Data handling · Dataset abstractions and feature conversion utilities. · `torch_rechub/data`
- [Utils](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/utils): Cross-cutting tools · Data helpers, model utilities, ONNX export, quantization, visualization, and task helpers. · `torch_rechub/utils`
- [Serving](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/serving): Retrieval serving · ANN and vector-store backends for deployment-oriented matching. · `torch_rechub/serving`
- [Ranking](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/ranking): CTR models · DeepFM, Wide&Deep, DCN, DIN, DIEN, BST, AFM, AutoInt, FiBiNET, and related ranking models. · `torch_rechub/models/ranking`
- [Matching](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/matching): Retrieval models · DSSM, YoutubeDNN, MIND, GRU4Rec, SASRec, NARM, STAMP, ComiRec, SINE, and related retrieval models. · `torch_rechub/models/matching`
- [Multi-task](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/multi_task): MTL models · ESMM, MMoE, PLE, AITM, and SharedBottom. · `torch_rechub/models/multi_task`
- [Generative](https://github.com/datawhalechina/torch-rechub/tree/main/torch_rechub/models/generative): Seq/gen models · HLLM, HSTU, TIGER, and RQVAE generative recommendation models. · `torch_rechub/models/generative`
- [CTR trainer](https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/ctr_trainer.py): Ranking trainer · Optimization and evaluation flow for ranking and CTR tasks. · `torch_rechub/trainers/ctr_trainer.py`
- [Match trainer](https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/match_trainer.py): Retrieval trainer · Training loop for two-tower and retrieval models. · `torch_rechub/trainers/match_trainer.py`
- [MTL trainer](https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/mtl_trainer.py): Multi-task trainer · Optimization flow for multi-task recommendation models. · `torch_rechub/trainers/mtl_trainer.py`
- [Seq trainer](https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/seq_trainer.py): Sequence trainer · Training loop for sequential recommendation and generative variants. · `torch_rechub/trainers/seq_trainer.py`
- [RQVAE trainer](https://github.com/datawhalechina/torch-rechub/blob/main/torch_rechub/trainers/rqvae_trainer.py): Generative trainer · Specialized trainer for RQVAE-style generative recommendation. · `torch_rechub/trainers/rqvae_trainer.py`

### Workflows

Runnable entry points for experiments, benchmarks, docs, and tutorials.

- [Benchmarks](https://github.com/datawhalechina/torch-rechub/tree/main/benchmarks): Benchmark harness · Config-driven suites, schemas, metrics, baselines, runners, and reports. · `benchmarks`
- [Examples](https://github.com/datawhalechina/torch-rechub/tree/main/examples): Runnable demos · Task-oriented scripts for ranking, matching, generative, and serving workflows. · `examples`
- [Docs](https://github.com/datawhalechina/torch-rechub/tree/main/docs): VitePress site · Bilingual documentation site for concepts, models, tools, serving, and tutorials. · `docs`
- [Tutorials](https://github.com/datawhalechina/torch-rechub/tree/main/tutorials): Notebooks · Interactive notebooks for quickstart, training, tracking, export, and serving. · `tutorials`

### Quality

Tests and repository automation that guard the framework contract.

- [Tests](https://github.com/datawhalechina/torch-rechub/tree/main/tests): E2E coverage · End-to-end checks for training, preprocessing, export, serving, and model behavior. · `tests`

## Connections

- torch_rechub → Basic: composes · Core package exposes shared primitives.
- torch_rechub → Models: contains · Core package defines model families.
- torch_rechub → Trainers: contains · Core package defines task trainers.
- torch_rechub → Data: contains · Core package includes data abstractions.
- torch_rechub → Utils: contains · Core package includes helpers and export tooling.
- torch_rechub → Serving: contains · Core package includes retrieval serving adapters.
- Models → Ranking: specializes · Ranking family implements CTR models.
- Models → Matching: specializes · Matching family implements retrieval models.
- Models → Multi-task: specializes · Multi-task family implements multi-objective models.
- Models → Generative: specializes · Generative family implements sequential/generative models.
- Ranking → Basic: reuses · Ranking models reuse shared layers, losses, and features.
- Matching → Basic: reuses · Matching models reuse shared layers, losses, and features.
- Multi-task → Basic: reuses · MTL models reuse shared layers, losses, and features.
- Generative → Basic: reuses · Generative models reuse shared primitives.
- Trainers → CTR trainer: includes · Ranking training path.
- Trainers → Match trainer: includes · Retrieval training path.
- Trainers → MTL trainer: includes · Multi-task training path.
- Trainers → Seq trainer: includes · Sequential training path.
- Trainers → RQVAE trainer: includes · Generative training path.
- Data → Trainers: feeds · Datasets and loaders feed task trainers.
- Examples → Models: demonstrates · Executable scripts show how to run models end to end.
- Examples → Trainers: drives · Examples call into task trainers.
- Examples → Serving: shows · Serving demos and deployment examples.
- Benchmarks → Models: evaluates · Benchmark suites compare model families.
- Benchmarks → Data: uses · Benchmark runs consume standardized datasets.
- Docs → Examples: documents · Docs explain runnable workflows and usage.
- Tutorials → Examples: teaches · Notebooks mirror the example pipelines.
- Tests → Trainers: verifies · E2E tests exercise task training loops.
- Tests → Serving: verifies · Tests cover export and serving behavior.
- Tests → Utils: verifies · Tests cover preprocessing, ONNX, and quantization utilities.