Welcome to deep-ml’s documentation!

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
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

Quick Start

Installation

pip install deepml

Basic Usage

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"
)

# 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",
    devices="auto"
)

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

Table of Contents

User Guide

API Reference

Additional Resources