Updated documentation to reflect new data loading workflow

Author: gperdrizet

docs/source/api/data.rst

@@ -9,66 +9,82 @@ Data loading
 Functions
 ---------
 
-.. autofunction:: image_classification_tools.pytorch.data.make_data_loaders
+.. autofunction:: image_classification_tools.pytorch.data.load_datasets
+
+.. autofunction:: image_classification_tools.pytorch.data.prepare_splits
+
+.. autofunction:: image_classification_tools.pytorch.data.create_dataloaders
+
+.. autofunction:: image_classification_tools.pytorch.data.generate_augmented_data
 
 Overview
 --------
 
-The data module provides flexible data loading capabilities with support for:
+The data module provides a flexible three-step data loading workflow:
+
+1. **Load datasets**: Load train/test datasets from PyTorch dataset classes or directories
+2. **Prepare splits**: Split data into train/val(/test) with configurable ratios
+3. **Create dataloaders**: Create DataLoaders with optional memory preloading strategies
 
-* torchvision datasets (CIFAR-10, MNIST, ImageFolder, etc.)
-* Custom train/eval transforms
-* Configurable batch sizes
-* Optional GPU preloading
-* Automatic train/validation splitting (default 80/20)
+Key features:
+
+* Support for torchvision datasets (CIFAR-10, MNIST, etc.) and custom ImageFolder datasets
+* Separate train and evaluation transforms
+* Flexible splitting: 2-way (train/val) or 3-way (train/val/test)
+* Three memory strategies: lazy loading, CPU preloading, or GPU preloading
+* Data augmentation with chunking for large datasets
+* Configurable batch sizes and workers
 
 Example usage
 -------------
 
-MNIST dataset:
+Basic workflow (CIFAR-10 with GPU preloading):
 
 .. code-block:: python
 
    from pathlib import Path
+   import torch
    from torchvision import datasets, transforms
-   from image_classification_tools.pytorch.data import make_data_loaders
+   from image_classification_tools.pytorch.data import (
+       load_datasets, prepare_splits, create_dataloaders
+   )
 
+   # Define transforms
    transform = transforms.Compose([
        transforms.ToTensor(),
-       transforms.Normalize((0.5,), (0.5,))
+       transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
 
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir=Path('./data'),
-       dataset_class=datasets.MNIST,
-       batch_size=128,
+   # Step 1: Load datasets
+   train_dataset, test_dataset = load_datasets(
+       data_source=datasets.CIFAR10,
        train_transform=transform,
-       eval_transform=transform
+       eval_transform=transform,
+       download=True,
+       root=Path('./data/cifar10')
    )
 
-CIFAR-10 dataset:
-
-.. code-block:: python
-
-   from torchvision import datasets
-
-   transform = transforms.Compose([
-       transforms.ToTensor(),
-       transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-   ])
+   # Step 2: Prepare splits (2-way: train/val from train_dataset)
+   train_dataset, val_dataset, test_dataset = prepare_splits(
+       train_dataset=train_dataset,
+       test_dataset=test_dataset,
+       train_val_split=0.8  # 80% train, 20% val
+   )
 
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir=Path('./data'),
-       dataset_class=datasets.CIFAR10,
+   # Step 3: Create dataloaders with GPU preloading
+   device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+   train_loader, val_loader, test_loader = create_dataloaders(
+       train_dataset, val_dataset, test_dataset,
        batch_size=128,
-       train_transform=transform,
-       eval_transform=transform
+       preload_to_memory=True,
+       device=device
    )
 
-With data augmentation:
+With data augmentation (lazy loading):
 
 .. code-block:: python
 
+   # Define separate transforms for training and evaluation
    train_transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
@@ -81,25 +97,48 @@ With data augmentation:
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
 
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir=Path('./data'),
-       dataset_class=datasets.CIFAR10,
-       batch_size=128,
+   # Load with different transforms
+   train_dataset, test_dataset = load_datasets(
+       data_source=datasets.CIFAR10,
        train_transform=train_transform,
        eval_transform=eval_transform,
-       device=None  # Keep on CPU for on-the-fly augmentation
+       root=Path('./data/cifar10')
    )
 
-Custom dataset with ImageFolder:
+   # Prepare splits
+   train_dataset, val_dataset, test_dataset = prepare_splits(
+       train_dataset=train_dataset,
+       test_dataset=test_dataset,
+       train_val_split=0.8
+   )
 
-.. code-block:: python
+   # Create dataloaders with lazy loading (no preloading)
+   train_loader, val_loader, test_loader = create_dataloaders(
+       train_dataset, val_dataset, test_dataset,
+       batch_size=128,
+       preload_to_memory=False,  # Lazy loading for augmentation
+       num_workers=4,
+       pin_memory=True
+   )
 
-   from torchvision.datasets import ImageFolder
+3-way split (no separate test set):
+
+.. code-block:: python
 
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir=Path('./my_dataset'),
-       dataset_class=ImageFolder,
-       batch_size=64,
+   # Load only training data (no test set available)
+   train_dataset, _ = load_datasets(
+       data_source=datasets.ImageFolder,
        train_transform=transform,
-       eval_transform=transform
+       eval_transform=transform,
+       root=Path('./my_dataset/train')
    )
+
+   # 3-way split: train/val/test all from train_dataset
+   train_dataset, val_dataset, test_dataset = prepare_splits(
+       train_dataset=train_dataset,
+       test_dataset=None,  # Will split test from train_dataset
+       train_val_split=0.8,  # 80/20 split of remaining data
+       test_split=0.15  # Reserve 15% for testing
+   )
+   # Results in approximately: 68% train, 17% val, 15% test
+

docs/source/index.rst

@@ -38,8 +38,11 @@ Minimal example classifying MNIST digits:
 .. code-block:: python
 
    import torch
+   from pathlib import Path
    from torchvision import datasets, transforms
-   from image_classification_tools.pytorch.data import make_data_loaders
+   from image_classification_tools.pytorch.data import (
+       load_datasets, prepare_splits, create_dataloaders
+   )
    from image_classification_tools.pytorch.training import train_model
 
    # Load data
@@ -48,12 +51,25 @@ Minimal example classifying MNIST digits:
        transforms.Normalize((0.5,), (0.5,))
    ])
    
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir='./data',
-       dataset_class=datasets.MNIST,
-       batch_size=64,
+   # Load, split, and create dataloaders
+   train_dataset, test_dataset = load_datasets(
+       data_source=datasets.MNIST,
        train_transform=transform,
-       eval_transform=transform
+       eval_transform=transform,
+       download=True,
+       root=Path('./data/mnist')
+   )
+   
+   train_dataset, val_dataset, test_dataset = prepare_splits(
+       train_dataset, test_dataset, train_val_split=0.8
+   )
+   
+   device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+   train_loader, val_loader, test_loader = create_dataloaders(
+       train_dataset, val_dataset, test_dataset,
+       batch_size=64,
+       preload_to_memory=True,
+       device=device
    )
 
    # Define model

docs/source/quickstart.rst

@@ -26,22 +26,39 @@ This example shows the complete workflow using the MNIST dataset.
    from pathlib import Path
    import torch
    from torchvision import datasets, transforms
-   from image_classification_tools.pytorch.data import make_data_loaders
+   from image_classification_tools.pytorch.data import (
+       load_datasets, prepare_splits, create_dataloaders
+   )
 
    # Define preprocessing
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
 
-   # Create data loaders
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir=Path('./data'),
-       dataset_class=datasets.MNIST,
-       batch_size=128,
+   # Step 1: Load datasets
+   train_dataset, test_dataset = load_datasets(
+       data_source=datasets.MNIST,
        train_transform=transform,
        eval_transform=transform,
-       device='cuda' if torch.cuda.is_available() else 'cpu'
+       download=True,
+       root=Path('./data/mnist')
+   )
+
+   # Step 2: Prepare splits
+   train_dataset, val_dataset, test_dataset = prepare_splits(
+       train_dataset=train_dataset,
+       test_dataset=test_dataset,
+       train_val_split=0.8
+   )
+
+   # Step 3: Create dataloaders
+   device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+   train_loader, val_loader, test_loader = create_dataloaders(
+       train_dataset, val_dataset, test_dataset,
+       batch_size=128,
+       preload_to_memory=True,
+       device=device
    )
 
 2. Define model
@@ -83,7 +100,7 @@ This example shows the complete workflow using the MNIST dataset.
        criterion=criterion,
        optimizer=optimizer,
        device=device,
-       lazy_loading=False,  # Data already on device from make_data_loaders
+       lazy_loading=False,  # Set to False when using preload_to_memory=True
        epochs=20,
        print_every=5
    )
@@ -124,16 +141,40 @@ For datasets in ImageFolder format:
 
 .. code-block:: python
 
+   from pathlib import Path
    from torchvision.datasets import ImageFolder
 
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir=Path('./my_dataset'),
-       dataset_class=ImageFolder,
-       batch_size=64,
+   # Define transform
+   transform = transforms.Compose([
+       transforms.Resize((224, 224)),
+       transforms.ToTensor(),
+       transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                          std=[0.229, 0.224, 0.225])
+   ])
+
+   # Load datasets from directory structure
+   train_dataset, test_dataset = load_datasets(
+       data_source=Path('./my_dataset'),
        train_transform=transform,
        eval_transform=transform
    )
 
+   # If no test directory exists, use 3-way split
+   train_dataset, val_dataset, test_dataset = prepare_splits(
+       train_dataset=train_dataset,
+       test_dataset=test_dataset,  # Will be None if no test/ directory
+       train_val_split=0.8,
+       test_split=0.1  # Only used if test_dataset is None
+   )
+
+   # Create dataloaders
+   train_loader, val_loader, test_loader = create_dataloaders(
+       train_dataset, val_dataset, test_dataset,
+       batch_size=64,
+       preload_to_memory=False,  # Lazy loading for large datasets
+       num_workers=4
+   )
+
 Your directory structure should be:
 
 .. code-block:: text
@@ -195,13 +236,26 @@ Improve generalization with data augmentation:
        transforms.Normalize((0.5,), (0.5,))
    ])
    
-   # Use different transforms for training and evaluation
-   train_loader, val_loader, test_loader = make_data_loaders(
-       data_dir=data_dir,
-       dataset_class=datasets.MNIST,
-       batch_size=128,
+   # Load with separate transforms
+   train_dataset, test_dataset = load_datasets(
+       data_source=datasets.MNIST,
        train_transform=train_transform,
-       eval_transform=eval_transform
+       eval_transform=eval_transform,
+       root=Path('./data/mnist')
+   )
+   
+   # Prepare splits
+   train_dataset, val_dataset, test_dataset = prepare_splits(
+       train_dataset, test_dataset, train_val_split=0.8
+   )
+   
+   # Create dataloaders with lazy loading (important for augmentation)
+   train_loader, val_loader, test_loader = create_dataloaders(
+       train_dataset, val_dataset, test_dataset,
+       batch_size=128,
+       preload_to_memory=False,  # Use lazy loading for on-the-fly augmentation
+       num_workers=4,
+       pin_memory=True
    )
 
 Hyperparameter optimization