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.6.tar.gz (182.2 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.6-py3-none-any.whl (250.5 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: deepsuite-1.0.6.tar.gz
  • Upload date:
  • Size: 182.2 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.6.tar.gz
Algorithm Hash digest
SHA256 5b6b781f843aa02641feca497e15d78cb153f6e88a6e52d8c4ea1786ed65b4c6
MD5 5b3243f3209ba6a59e3492548cfa75e4
BLAKE2b-256 acdf460b4f8ac6d2fe8d02493dae4f22b8d0a5fc4f4552425fb3af70e5d69802

See more details on using hashes here.

File details

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

File metadata

  • Download URL: deepsuite-1.0.6-py3-none-any.whl
  • Upload date:
  • Size: 250.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.6-py3-none-any.whl
Algorithm Hash digest
SHA256 4611bbe256d66b8c2e267670600d678b5c716dccb8411609f8bb1a02c7b59fb2
MD5 49ebffd6a0434464eb399d86d313024b
BLAKE2b-256 20d466397a4d5dc68f4083b2b943c53094db0d0f6dc5baead4b5c9880080da7a

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