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)

Feature Matching (LoFTR)

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.3 | 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.0.tar.gz (161.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.0-py3-none-any.whl (224.0 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: deepsuite-1.0.0.tar.gz
  • Upload date:
  • Size: 161.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.0.tar.gz
Algorithm Hash digest
SHA256 b692c7b6ab85a87e22d591aab4e4bb1a8d8e0771e21fcf1aacdcab9b5481d66b
MD5 00ae5862c972461f9e594471c77f546e
BLAKE2b-256 cb3e321a5de57aed32b31003becf2eb344ae9fd38a6664d1b4cdcd933841dbfe

See more details on using hashes here.

File details

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

File metadata

  • Download URL: deepsuite-1.0.0-py3-none-any.whl
  • Upload date:
  • Size: 224.0 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.0-py3-none-any.whl
Algorithm Hash digest
SHA256 8b696080a89395a24fa096985e58aa434431461ffb48d84f75145001a16410d8
MD5 9a7c1bef637fc704d96121803f2d98f1
BLAKE2b-256 6e4b2ed06a709016c2400988de78b53fccd52b13a088a748f0f864993bfe31e0

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