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 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
- LLM Modules: GPT and DeepSeek-V3 architecture documentation
- LLM Loss & Heads: Language modeling losses and output heads
- Mixture-of-Experts: DeepSeek-V3 MoE implementation
- Text Datasets: Text data loading and preprocessing
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:
- Fork the repository
- Create a feature branch (`git checkout -b feature/amazing-feature`)
- Commit your changes (`git commit -m 'Add amazing feature'`)
- Push to the branch (`git push origin feature/amazing-feature`)
- Open a Pull Request
Please ensure:
- Code follows the style guide (Ruff + MyPy)
- Tests pass (pytest)
- Documentation is updated
๐ Acknowledgments
- PyTorch Lightning for the training framework
- DeepSeek-AI for the DeepSeek-V3 architecture
- Ultralytics for YOLO implementations
- The open-source community for various model implementations
๐ง 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
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
f21a6ea09394765eabd16c8c77743659ac65960b47c2fe7aee8ef350e30c6e4a
|
|
| MD5 |
b7cff361ada01d83c9ac04d2e58104b7
|
|
| BLAKE2b-256 |
6b5b21decc610933bda14c57cb3348c575d18e0546f20243a23111436623d40d
|
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
a307fc40a7464ee156350dfb6c27814cacd9aacbd4d6ae60bd05a81f95f69075
|
|
| MD5 |
6600c46283bbaa24741f2c35c4fe5c63
|
|
| BLAKE2b-256 |
31587c998955bda639af4248e3c2b6f541358eaf9962af1f350e6dc9a0f77ccb
|