deep-ml

deep-ml is a high-level PyTorch training framework that simplifies deep learning workflows for computer vision tasks. It provides easy-to-use trainers with distributed training support, comprehensive task implementations, and seamless experiment tracking.

Key Features

Multiple Training Backends

• FabricTrainer: Lightning Fabric for distributed training (recommended for multi-GPU)
• AcceleratorTrainer: HuggingFace Accelerate integration (recommended for multi-GPU)
• Learner: Classic PyTorch trainer (single-device, notebook-friendly)

Pre-built Task Implementations

• Image Classification (single & multi-label)
• Semantic Segmentation (binary & multiclass)
• Image Regression
• Custom tasks via extensible base classes

Experiment Tracking

• TensorBoard integration (default)
• MLflow support
• Weights & Biases (wandb) integration
• Custom logger interface

Advanced Training Features

• ✅ Automatic Mixed Precision (AMP)
• ✅ Gradient accumulation & clipping
• ✅ Learning rate scheduling with warmup
• ✅ Multi-GPU and distributed training
• ✅ Checkpoint management
• ✅ Progress bars and real-time metrics

Installation

Basic Installation

pip install deepml

With Optional Dependencies

# For Lightning Fabric
pip install deepml lightning-fabric

# For HuggingFace Accelerate
pip install deepml accelerate

# For MLflow tracking
pip install deepml mlflow

# For Weights & Biases
pip install deepml wandb

# For Albumentations (segmentation)
pip install deepml albumentations

Quick Start

Image Classification

from deepml.tasks import ImageClassification
from deepml.fabric_trainer import FabricTrainer
import torch
from torch.optim import Adam
from torchvision.models import resnet18

# 1. Define your model
model = resnet18(num_classes=10)

# 2. Create a task
task = ImageClassification(
    model=model,
    model_dir="./checkpoints",
    classes=['cat', 'dog', 'bird', ...]  # Optional
)

# 3. Setup optimizer and loss
optimizer = Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.CrossEntropyLoss()

# 4. Create trainer
trainer = FabricTrainer(
    task=task,
    optimizer=optimizer,
    criterion=criterion,
    accelerator="auto",  # Use GPU if available
    devices="auto",      # Use all available devices
    precision="16-mixed" # Mixed precision training
)

# 5. Train!
trainer.fit(
    train_loader=train_loader,
    val_loader=val_loader,
    epochs=50
)

# 6. Visualize predictions
task.show_predictions(loader=val_loader, samples=9)

Semantic Segmentation

from deepml.tasks import Segmentation
from deepml.fabric_trainer import FabricTrainer
from deepml.losses import JaccardLoss

# Define model (e.g., U-Net)
model = UNet(in_channels=3, out_channels=1)

# Create task
task = Segmentation(
    model=model,
    model_dir="./checkpoints",
    mode="binary",
    num_classes=1,
    threshold=0.5
)

# Setup training
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = torch.nn.BCEWithLogitsLoss()

trainer = FabricTrainer(task=task, optimizer=optimizer, criterion=criterion)

# Train
trainer.fit(
    train_loader=train_loader,
    val_loader=val_loader,
    epochs=100
)

Documentation

📚 Full documentation is available at: https://deep-ml.readthedocs.io/

Documentation Structure

Getting Started

• Installation Guide
• Quick Start Tutorial
• Tutorials

User Guide

• Trainers - FabricTrainer, AcceleratorTrainer, Learner
• Tasks - Classification, Segmentation, Regression
• Datasets - Data loading utilities
• Loss Functions - Custom losses for CV tasks
• Metrics - Evaluation metrics
• Tracking - MLflow, TensorBoard, Wandb
• Visualization - Result visualization

API Reference

• Complete API Documentation

Additional Resources

• Examples
• FAQ
• Changelog
• Contributing

Tutorials

Available Tutorials

1. Image Classification: Train ResNet on CIFAR-10
2. Transfer Learning: Fine-tune pre-trained models
3. Semantic Segmentation: U-Net for binary segmentation
4. Multi-GPU Training: Distributed training across GPUs
5. Hyperparameter Tuning: Optimize with Optuna
6. Model Deployment: Export to TorchScript/ONNX

📖 See the complete tutorials on ReadTheDocs.

💡 Advanced Features

Distributed Training

# Multi-GPU training with DDP
trainer = FabricTrainer(
    task=task,
    optimizer=optimizer,
    criterion=criterion,
    accelerator="gpu",
    strategy="ddp",
    devices="auto"  # Use all GPUs
)

Gradient Accumulation

# Simulate larger batch sizes
trainer.fit(
    train_loader=train_loader,
    val_loader=val_loader,
    epochs=50,
    gradient_accumulation_steps=4  # Effective batch = 4x
)

Learning Rate Scheduling

from deepml.lr_scheduler_utils import setup_one_cycle_lr_scheduler_with_warmup

lr_scheduler_fn = lambda opt: setup_one_cycle_lr_scheduler_with_warmup(
    optimizer=opt,
    steps_per_epoch=len(train_loader),
    warmup_ratio=0.1,
    num_epochs=50,
    max_lr=1e-3
)

trainer = FabricTrainer(
    ...,
    lr_scheduler_fn=lr_scheduler_fn
)

steps_per_epoch counts optimizer steps, so len(train_loader) is only correct for a single process with no gradient accumulation. Accumulation reduces the count, and Fabric shards the loader inside fit() — after the scheduler is built — so account for both yourself:

import math

batches_per_rank = math.ceil(len(train_loader) / num_processes)
steps_per_epoch = math.ceil(batches_per_rank / gradient_accumulation_steps)

Use ceiling division and prefer over-estimating: too large only stops the cycle short of its final min LR, while too small makes OneCycleLR raise ValueError once the trainer steps past total_steps.

Experiment Tracking

from deepml.tracking import MLFlowLogger, WandbLogger

# MLflow
logger = MLFlowLogger(
    experiment_name='my-experiment',
    tracking_uri='./mlruns'
)

# Weights & Biases
logger = WandbLogger(
    project='my-project',
    name='experiment-1'
)

trainer.fit(..., logger=logger)

Supported Tasks

Task	Description	Typical Use Cases
ImageClassification	Single-label classification	CIFAR-10, ImageNet
MultiLabelImageClassification	Multi-label classification	Object attributes
Segmentation	Pixel-level classification	Medical imaging, autonomous driving
ImageRegression	Continuous value prediction	Age estimation, depth prediction
NeuralNetTask	Generic task template	Custom tasks

Custom Loss Functions

• JaccardLoss: IoU loss for segmentation
• RMSELoss: Root mean squared error
• WeightedBCEWithLogitsLoss: Weighted binary cross-entropy
• ContrastiveLoss: For siamese networks
• AngularPenaltySMLoss: ArcFace, SphereFace, CosFace for face recognition

Metrics

• Classification: Accuracy, BinaryAccuracy
• Segmentation: IoU, Dice Coefficient, Pixel Accuracy
• Custom: Easy to implement custom metrics

Datasets

• ImageDataFrameDataset: Load from pandas DataFrame
• ImageRowDataFrameDataset: Flattened arrays in DataFrame
• SegmentationDataFrameDataset: Images + masks with Albumentations
• ImageListDataset: Directory of images

Contributing

Contributions are welcome! See our Contributing Guide for guidelines.

Development Setup

git clone https://github.com/sagar100rathod/deep-ml.git
cd deep-ml
pip install -e ".[dev]"
pytest  # Run tests

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• PyTorch team for the amazing framework
• Lightning AI for Lightning Fabric
• HuggingFace for Accelerate
• All contributors to this project

Contact

• Author: Sagar Rathod
• Email: sagar100rathod@gmail.com
• Issues: GitHub Issues
• Discussions: GitHub Discussions

⭐ Star History

If you find this project useful, please consider giving it a star!

Citation

If you use deep-ml in your research, please cite:

@software{deepml2026,
  author    = {Sagar Rathod},
  title     = {deep-ml: A High-Level PyTorch Training Framework for Computer Vision},
  year      = {2026},
  publisher = {GitHub},
  url       = {https://github.com/sagar100rathod/deep-ml},
  doi       = {10.5281/zenodo.1234567}
}