Skip to main content

Comprehensive deep learning framework with state-of-the-art implementations: GPT, DeepSeek-V3 with MLA/MoE, YOLO, CenterNet, and specialized audio/vision models built on PyTorch Lightning

Project description

DeepSuite

DeepSuite

DeepSuite is a comprehensive deep learning framework based on PyTorch Lightning. It provides production-ready implementations of modern architectures: language models (GPT, DeepSeekโ€‘V3 with MLA/MoE), object detection (YOLO, CenterNet), and specialized audio/vision models. Docstrings and examples follow Googleโ€‘Style.

๐Ÿš€ Features

Language Models & NLP

  • โœ… GPT-2/GPT-3 Architecture: Full transformer implementation with configurable layers
  • โœ… DeepSeek-V3: State-of-the-art LLM with Multi-Head Latent Attention (MLA) and Mixture-of-Experts (MoE)
    • Multi-Head Latent Attention with KV-Compression
    • Auxiliary-Loss-Free Load Balancing
    • Multi-Token Prediction (MTP)
    • Rotary Position Embeddings (RoPE)
  • โœ… Text Dataset Loaders: Support for .txt, .jsonl, pre-tokenized data with sliding window

Computer Vision

  • โœ… Object Detection: YOLO (v3/v4/v5), CenterNet, EfficientDet
  • โœ… Feature Extractors: ResNet, DarkNet, EfficientNet, MobileNet, FPN
  • โœ… Tracking: RNN-based object tracking with ReID

Audio Processing

  • โœ… Beamforming: Multi-channel audio processing
  • โœ… Direction of Arrival (DOA): Acoustic source localization
  • โœ… Feature Networks: WaveNet, Complex-valued networks

Training & Optimization

  • ๐Ÿง  Continual Learning: Built-in knowledge distillation and LwF (Learning without Forgetting)
  • ๐Ÿงฉ Pluggable Callbacks: TorchScript export, TensorRT optimization, t-SNE visualization
  • โš™๏ธ Modular Design: Composable heads, losses, layers, and metrics
  • ๐Ÿ“Š Visualization Tools: Embedding analysis, training metrics, model profiling
  • ๐Ÿ—‚๏ธ Flexible Dataset Loaders: Image, audio, text with augmentation support

๐Ÿ› ๏ธ Installation

# Clone repository
git clone https://github.com/afeldman/deepsuite.git
cd deepsuite

# Virtual environment (recommended)
python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate

# Installation (CPU)
uv sync --extra cpu --extra dev

# Installation (CUDA 12.4)
uv sync --extra cu124 --extra dev

# Documentation build dependencies
brew install graphviz             # macOS
sudo apt-get install libgraphviz-dev  # Linux

๐ŸŽฏ Quick Start

Language Model Training (DeepSeekโ€‘V3)

import pytorch_lightning as pl
from deepsuite.model.llm.deepseek import DeepSeekV3Module
from deepsuite.lightning_base.dataset.text_loader import TextDataLoader

# Prepare dataset
datamodule = TextDataLoader(
    train_data_path="data/train.txt",
    val_data_path="data/val.txt",
    tokenizer=tokenizer,
    max_seq_len=512,
    batch_size=32,
    return_mtp=True,  # Enable Multi-Token Prediction
)

# Create model
model = DeepSeekV3Module(
    vocab_size=50000,
    d_model=2048,
    n_layers=24,
    n_heads=16,
    use_moe=True,           # Mixture-of-Experts
    use_mtp=True,           # Multi-Token Prediction
    n_routed_experts=256,
    n_expert_per_token=8,
)

# Training
trainer = pl.Trainer(max_steps=100000, accelerator="gpu")
trainer.fit(model, datamodule)

Object Detection (YOLO)

import pytorch_lightning as pl
from deepsuite.model.feature.yolo import YOLOLightningModule

model = YOLOLightningModule(
    num_classes=80,
    backbone="cspdarknet",
    use_rotated_loss=False,
)

trainer = pl.Trainer(max_epochs=100, accelerator="gpu")
trainer.fit(model, datamodule)

### Pluggable Heads & Global Registry

DeepSuite provides a global head registry and trainer helpers to compose multi-head models across projects.

```python
from deepsuite.registry import HeadRegistry
from deepsuite.lightning_base.trainer import train_heads_with_registry

# Register your head class in your project
@HeadRegistry.register("my_head")
class MyHead:
    ...

# Build and train via registry
trainer = train_heads_with_registry(
    head_cfgs=[{"name": "my_head", "args": {"param": 123}}],
    module_builder=MyMultiHeadLightningModule,  # accepts heads, share_backbone
    datamodule=my_data_module,
    trainer_cfg={"max_epochs": 10, "accelerator": "gpu"},
    share_backbone=True,
)

### Feature Matching (LoFTR)

```python
from deepsuite.model.loftr.loftr import LoFTR

loftr = LoFTR(d_model=256, nhead=8)
matches = loftr(img1, img2)

Spatial Transformer Networks (STN)

from deepsuite.model.stn import AffineSTN

stn = AffineSTN(in_channels=3)
warped = stn(images)

๐Ÿ“š Documentation

Core Modules

Language Models

Computer Vision

  • Detection models (YOLO, CenterNet, EfficientDet)
  • Feature extractors (ResNet, DarkNet, EfficientNet, FPN)
  • Object tracking systems

Audio Processing

  • Beamforming algorithms
  • DOA estimation
  • Complex-valued neural networks

Examples

# Language Models
python examples/llm_modules_example.py      # GPT & DeepSeek-V3
python examples/llm_loss_head_example.py    # Loss functions & heads
python examples/moe_example.py              # Mixture-of-Experts
python examples/text_dataset_simple.py      # Text data loading

# View all examples
ls examples/

Build Documentation Locally

cd docs
make html
# Open docs/_build/html/index.html

๐Ÿ—๏ธ Project Structure

deepsuite/
โ”œโ”€โ”€ src/deepsuite/                      # Core source code
โ”‚   โ”œโ”€โ”€ callbacks/                 # Training callbacks (TensorRT, t-SNE, etc.)
โ”‚   โ”œโ”€โ”€ heads/                     # Output heads (classification, detection, LM)
โ”‚   โ”œโ”€โ”€ layers/                    # Neural network layers
โ”‚   โ”‚   โ”œโ”€โ”€ attention/             # Attention mechanisms (MLA, RoPE, KV-Compression)
โ”‚   โ”‚   โ””โ”€โ”€ moe.py                 # Mixture-of-Experts
โ”‚   โ”œโ”€โ”€ lightning_base/            # Lightning modules and utilities
โ”‚   โ”‚   โ””โ”€โ”€ dataset/               # Dataset loaders (image, audio, text)
โ”‚   โ”œโ”€โ”€ loss/                      # Loss functions
โ”‚   โ”œโ”€โ”€ metric/                    # Evaluation metrics
โ”‚   โ”œโ”€โ”€ model/                     # Model architectures
โ”‚   โ”‚   โ”œโ”€โ”€ beamforming/           # Audio beamforming
โ”‚   โ”‚   โ””โ”€โ”€ detection/             # Object detection
โ”‚   โ”œโ”€โ”€ modules/                   # Lightning modules
โ”‚   โ”‚   โ”œโ”€โ”€ deepseek.py            # DeepSeek-V3 module
โ”‚   โ”‚   โ”œโ”€โ”€ gpt.py                 # GPT module
โ”‚   โ”‚   โ”œโ”€โ”€ yolo.py                # YOLO module
โ”‚   โ”‚   โ””โ”€โ”€ ...
โ”‚   โ””โ”€โ”€ utils/                     # Utility functions
โ”œโ”€โ”€ examples/                      # Usage examples
โ”‚   โ”œโ”€โ”€ llm_modules_example.py     # Language model examples
โ”‚   โ”œโ”€โ”€ moe_example.py             # MoE examples
โ”‚   โ””โ”€โ”€ text_dataset_simple.py     # Dataset examples
โ”œโ”€โ”€ tests/                         # Unit tests
โ”œโ”€โ”€ docs/                          # Sphinx documentation
โ”‚   โ”œโ”€โ”€ llm_modules.md             # LLM documentation
โ”‚   โ”œโ”€โ”€ moe.md                     # MoE documentation
โ”‚   โ””โ”€โ”€ text_dataset.md            # Dataset documentation
โ”œโ”€โ”€ pyproject.toml                 # Project configuration
โ”œโ”€โ”€ README.md                      # This file
#

๐ŸŽ“ Key Concepts

Multi-Head Latent Attention (MLA)

DeepSeek-V3's efficient attention mechanism with low-rank KV-compression:

from deepsuite.layers.attention.mla import MultiHeadLatentAttention

attention = MultiHeadLatentAttention(
    d=2048,  # Model dimension
    n_h=16,  # Number of heads
    d_h_c=256,  # Compressed KV dimension
    d_h_r=64,  # Per-head RoPE dimension
)

Mixture-of-Experts (MoE)

Sparse expert activation with auxiliary-loss-free load balancing:

from deepsuite.layers.moe import DeepSeekMoE

moe = DeepSeekMoE(
    d_model=2048,
    n_shared_experts=1,  # Always active
    n_routed_experts=256,  # Selectively activated
    n_expert_per_token=8,  # Top-K experts per token
)

Multi-Token Prediction (MTP)

Predict multiple future tokens simultaneously:

# Dataset mit MTP-Zielen
dataset = TextDataset(
    data_path="train.txt",
    tokenizer=tokenizer,
    return_mtp=True,
    mtp_depth=3,  # Predict 1, 2, 3 tokens ahead
)

# Model with MTP loss
model = DeepSeekV3Module(
    vocab_size=50000,
    use_mtp=True,
    mtp_lambda=0.3,  # MTP loss weight
)

๐Ÿ“Š Model Zoo

Language Models

Model Parameters Config Performance
GPT-Small 124M `d_model=768, n_layers=12` GPT-2 baseline
DeepSeek-Small 51M `d_model=512, n_layers=4, use_moe=True` Demo config
DeepSeek-Base 1.3B `d_model=2048, n_layers=24, n_experts=256` Production
DeepSeek-V3 685B `d_model=7168, n_layers=60, n_experts=256` Full scale

Object Detection

Model Backbone mAP FPS
YOLOv5s CSPDarknet 37.4 140
YOLOv5m CSPDarknet 45.4 100
CenterNet ResNet-50 42.1 45

๐Ÿงช Testing

# All tests
pytest

# Single test
pytest tests/test_tensor_rt_export_callback.py

# With coverage
pytest --cov=deepsuite

๐Ÿ› ๏ธ Development

Code Quality (Googleโ€‘Style, Ruff, MyPy)

Docstrings follow Google-Style and are verified via Ruff (pydocstyle=google).

# Format code
ruff format .

# Linting (Auto-Fix)
ruff check . --fix

# Type checking
mypy src/deepsuite

Optional: Set up pre-commit hooks.

# Format code
ruff format .

# Lint
ruff check .

# Type checking
mypy src/deepsuite

Preโ€‘commit Hooks

# Install pre-commit
pip install pre-commit

# Setup hooks
pre-commit install

# Run manually
pre-commit run --all-files

๐Ÿ“– Citation

If you use DeepSuite in your research, please cite:

@software{deepsuite2025,
title = {DeepSuite},
author = {Anton Feldmann},
year = {2025},
url = {https://github.com/afeldman/deepsuite}
}

For DeepSeek-V3:

@article{deepseekai2024deepseekv3,
title={DeepSeek-V3 Technical Report},
author={DeepSeek-AI},
journal={arXiv preprint arXiv:2412.19437},
year={2024}
}

๐Ÿค Contributing

Contributions are welcome! Please:

  1. Fork the repository
  2. Create a feature branch (`git checkout -b feature/amazing-feature`)
  3. Commit your changes (`git commit -m 'Add amazing feature'`)
  4. Push to the branch (`git push origin feature/amazing-feature`)
  5. Open a Pull Request

Please ensure:

  • Code follows the style guide (Ruff + MyPy)
  • Tests pass (pytest)
  • Documentation is updated

๐Ÿ™ Acknowledgments

๐Ÿ“ง Contact

Anton Feldmann - anton.feldmann@gmail.com

Project Link: https://github.com/afeldman/deepsuite


Version: 1.0.1 | Python: >=3.11 | PyTorch: >=2.6.0 | Lightning: >=2.5.1 | License: Apache 2.0

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

deepsuite-1.0.4.tar.gz (166.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

deepsuite-1.0.4-py3-none-any.whl (230.5 kB view details)

Uploaded Python 3

File details

Details for the file deepsuite-1.0.4.tar.gz.

File metadata

  • Download URL: deepsuite-1.0.4.tar.gz
  • Upload date:
  • Size: 166.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.2

File hashes

Hashes for deepsuite-1.0.4.tar.gz
Algorithm Hash digest
SHA256 f21a6ea09394765eabd16c8c77743659ac65960b47c2fe7aee8ef350e30c6e4a
MD5 b7cff361ada01d83c9ac04d2e58104b7
BLAKE2b-256 6b5b21decc610933bda14c57cb3348c575d18e0546f20243a23111436623d40d

See more details on using hashes here.

File details

Details for the file deepsuite-1.0.4-py3-none-any.whl.

File metadata

  • Download URL: deepsuite-1.0.4-py3-none-any.whl
  • Upload date:
  • Size: 230.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.2

File hashes

Hashes for deepsuite-1.0.4-py3-none-any.whl
Algorithm Hash digest
SHA256 a307fc40a7464ee156350dfb6c27814cacd9aacbd4d6ae60bd05a81f95f69075
MD5 6600c46283bbaa24741f2c35c4fe5c63
BLAKE2b-256 31587c998955bda639af4248e3c2b6f541358eaf9962af1f350e6dc9a0f77ccb

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page