Datasets et dataLoaders PyTorch : Gérer les données efficacement

11 min de lecture 2274 mots

tl;dr: Dataset = accès données (__len__, __getitem__). DataLoader = batching/shuffle/parallélisation (num_workers). Transforms : RandomCrop, RandomFlip, ColorJitter, Normalize. Classes déséquilibrées : WeightedRandomSampler. Optimisations : pin_memory=True (GPU), persistent_workers=True. Héritage Dataset pour custom data.

Maintenant que vous savez entraîner un modèle, il est temps d’apprendre à gérer vos données efficacement ! Dans le deep learning, la qualité et la préparation des données sont aussi importantes que l’architecture du modèle. PyTorch fournit des outils puissants pour charger, transformer et augmenter vos données : Dataset et DataLoader.

Vue d’ensemble : Dataset et DataLoader

PyTorch sépare la gestion des données en deux composants :

Dataset : Définit comment accéder aux données (lecture, prétraitement)
DataLoader : Gère le batching, shuffling et parallélisation

from torch.utils.data import Dataset, DataLoader

# Dataset : "Quelles sont mes données et comment y accéder ?"
dataset = MyCustomDataset(data_path)

# DataLoader : "Comment créer des batches et les charger efficacement ?"
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

# Utilisation dans training loop
for batch_data, batch_labels in dataloader:
    # batch_data : (32, ...) - 32 exemples
    # Entraîner le modèle...
    pass

Schéma Dataset et DataLoader avec pipeline de données

La classe dataset : Accéder aux données

Dataset est une classe abstraite qui définit comment accéder à vos données. Vous devez implémenter deux méthodes :

Structure de base d’un dataset

from torch.utils.data import Dataset
import torch

class CustomDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        """
        Initialisation : charger les chemins de fichiers, métadonnées, etc.
        """
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        """
        Retourne la taille du dataset
        """
        return len(self.data)

    def __getitem__(self, idx):
        """
        Charge et retourne un exemple à l'index idx
        """
        # Charger l'exemple
        x = self.data[idx]
        y = self.labels[idx]

        # Appliquer transformations si spécifiées
        if self.transform:
            x = self.transform(x)

        return x, y

# Utilisation
data = torch.randn(100, 3, 32, 32)  # 100 images 32x32 RGB
labels = torch.randint(0, 10, (100,))  # 100 labels (10 classes)

dataset = CustomDataset(data, labels)
print(f"Dataset size: {len(dataset)}")  # 100

# Accéder à un exemple
x, y = dataset[0]
print(f"First sample: x.shape={x.shape}, label={y}")

💡 Les 3 méthodes essentielles : __init__() pour initialiser (lire chemins de fichiers, métadonnées), __len__() pour retourner la taille, et __getitem__(idx) pour charger un exemple spécifique. PyTorch appelle ces méthodes automatiquement.

Dataset pour images depuis dossiers

Cas d’usage très courant : images organisées en dossiers par classe

from torch.utils.data import Dataset
from PIL import Image
import os

class ImageFolderDataset(Dataset):
    """
    Structure attendue :
    root_dir/
        class_0/
            img1.jpg
            img2.jpg
        class_1/
            img3.jpg
            img4.jpg
    """
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.images = []
        self.labels = []

        # Parcourir les dossiers (un dossier = une classe)
        for label, class_name in enumerate(sorted(os.listdir(root_dir))):
            class_path = os.path.join(root_dir, class_name)
            if not os.path.isdir(class_path):
                continue

            # Lister les images dans ce dossier
            for img_name in os.listdir(class_path):
                if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                    img_path = os.path.join(class_path, img_name)
                    self.images.append(img_path)
                    self.labels.append(label)

        print(f"Loaded {len(self.images)} images from {len(set(self.labels))} classes")

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        # Charger l'image
        img_path = self.images[idx]
        image = Image.open(img_path).convert('RGB')

        # Appliquer transformations
        if self.transform:
            image = self.transform(image)

        label = self.labels[idx]
        return image, label

# Utilisation
# dataset = ImageFolderDataset('data/train/', transform=transforms.ToTensor())

Dataset depuis CSV

Pour des données tabulaires ou des chemins d’images dans un CSV :

import pandas as pd
from torch.utils.data import Dataset
from PIL import Image

class CSVImageDataset(Dataset):
    """
    CSV avec colonnes : image_path, label
    """
    def __init__(self, csv_file, root_dir, transform=None):
        self.data_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.data_frame)

    def __getitem__(self, idx):
        # Lire la ligne du CSV
        row = self.data_frame.iloc[idx]

        # Charger l'image
        img_name = os.path.join(self.root_dir, row['image_path'])
        image = Image.open(img_name).convert('RGB')

        # Récupérer le label
        label = int(row['label'])

        # Transformations
        if self.transform:
            image = self.transform(image)

        return image, label

# Utilisation
# dataset = CSVImageDataset('labels.csv', 'data/images/', transform=transforms.ToTensor())

Dataset en mémoire (petit dataset)

Pour des datasets qui tiennent en RAM :

class InMemoryDataset(Dataset):
    def __init__(self, data, labels):
        """
        data: tensor de shape (N, ...)
        labels: tensor de shape (N,)
        """
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# Exemple avec données synthétiques
data = torch.randn(1000, 3, 32, 32)
labels = torch.randint(0, 10, (1000,))
dataset = InMemoryDataset(data, labels)

Transformations avec torchvision.transforms

torchvision.transforms fournit des transformations courantes pour images.

Transformations de base

from torchvision import transforms

# Composer plusieurs transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),        # Redimensionner
    transforms.ToTensor(),                 # PIL Image → Tensor [0, 1]
    transforms.Normalize(                  # Normaliser
        mean=[0.485, 0.456, 0.406],       # ImageNet mean
        std=[0.229, 0.224, 0.225]          # ImageNet std
    )
])

# Appliquer
from PIL import Image
img = Image.open('image.jpg')
img_tensor = transform(img)
print(f"Transformed: {img_tensor.shape}")  # (3, 224, 224)

Transformations courantes

# Redimensionnement
transforms.Resize(256)                    # Plus petit côté = 256
transforms.Resize((224, 224))             # Exact 224x224
transforms.CenterCrop(224)                # Crop central 224x224

# Conversion
transforms.ToTensor()                     # PIL → Tensor [0, 1]
transforms.ToPILImage()                   # Tensor → PIL

# Normalisation
transforms.Normalize(mean, std)           # (x - mean) / std

# Flip
transforms.RandomHorizontalFlip(p=0.5)    # Flip horizontal 50%
transforms.RandomVerticalFlip(p=0.5)      # Flip vertical 50%

# Rotation
transforms.RandomRotation(degrees=15)     # Rotation ±15°

# Color
transforms.ColorJitter(
    brightness=0.2,    # ±20% luminosité
    contrast=0.2,      # ±20% contraste
    saturation=0.2,    # ±20% saturation
    hue=0.1           # ±10% teinte
)

# Grayscale
transforms.Grayscale(num_output_channels=3)  # RGB → Grayscale (3 canaux)

Data augmentation : Augmenter artificiellement les données

La data augmentation crée des variantes des données d’entraînement pour améliorer la généralisation.

Transformations d’entraînement vs validation

from torchvision import transforms

# Transformations TRAINING (avec augmentation)
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),  # Crop aléatoire
    transforms.RandomHorizontalFlip(p=0.5),                # Flip 50%
    transforms.ColorJitter(
        brightness=0.2,
        contrast=0.2,
        saturation=0.2,
        hue=0.1
    ),
    transforms.RandomRotation(15),                         # Rotation ±15°
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# Transformations VALIDATION (sans augmentation!)
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# Utiliser des transforms différents
train_dataset = ImageFolderDataset('data/train/', transform=train_transform)
val_dataset = ImageFolderDataset('data/val/', transform=val_transform)

⚠️ Warning
Crucial : N’appliquez jamais de data augmentation sur les données de validation/test ! Cela fausserait les métriques. Validation = transformations déterministes seulement (resize, crop, normalize).

Augmentation avancée avec RandAugment

from torchvision.transforms import RandAugment

# RandAugment : augmentation aléatoire automatique
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    RandAugment(num_ops=2, magnitude=9),  # Applique 2 ops aléatoires
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

AutoAugment (policies pré-apprises)

from torchvision.transforms import AutoAugment, AutoAugmentPolicy

# Policies optimisées pour ImageNet, CIFAR10, SVHN
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    AutoAugment(policy=AutoAugmentPolicy.IMAGENET),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

DataLoader : Batching et parallélisation

DataLoader prend un Dataset et gère le batching, shuffling, et chargement parallèle.

Paramètres principaux

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=32,           # Taille des batches
    shuffle=True,            # Mélanger les données
    num_workers=4,           # Nombre de processus pour chargement
    pin_memory=True,         # Optimisation GPU
    drop_last=True,          # Ignorer dernier batch incomplet
    persistent_workers=True  # Garder workers entre epochs
)

# Itérer sur les batches
for batch_idx, (data, targets) in enumerate(dataloader):
    print(f"Batch {batch_idx}: data.shape={data.shape}, targets.shape={targets.shape}")
    # data: (32, 3, 224, 224)
    # targets: (32,)

Comprendre num_workers

num_workers contrôle le nombre de processus qui chargent les données en parallèle :

import time
from torch.utils.data import DataLoader

dataset = MyDataset(...)  # Dataset volumineux

# Test avec différents num_workers
for num_workers in [0, 2, 4, 8]:
    dataloader = DataLoader(
        dataset,
        batch_size=32,
        num_workers=num_workers
    )

    start = time.time()
    for batch in dataloader:
        pass  # Juste charger les données
    elapsed = time.time() - start

    print(f"num_workers={num_workers}: {elapsed:.2f}s")

# Résultats typiques :
# num_workers=0: 15.23s  (séquentiel)
# num_workers=2: 8.45s   (1.8x faster)
# num_workers=4: 5.12s   (3.0x faster)
# num_workers=8: 4.89s   (3.1x faster, diminishing returns)

🔎 Tip
Choisir num_workers : Commencez avec num_workers = 4 (ou 2×nombre de GPUs). Trop de workers = overhead. Optimal dépend de : vitesse CPU, vitesse disque, complexité transforms. Testez 2, 4, 8 et mesurez.

Optimisations DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,

    # GPU optimizations
    pin_memory=True,          # Alloue en pinned memory (CPU → GPU plus rapide)

    # Worker optimizations
    persistent_workers=True,  # Garde workers entre epochs (évite reload)
    prefetch_factor=2,        # Nombre de batches pré-chargés par worker

    # Memory optimization
    drop_last=True            # Ignore dernier batch si incomplet
)

Collate function personnalisée

Pour gérer des données de tailles variables :

def custom_collate(batch):
    """
    batch: liste de tuples (data, label)
    """
    # Séparer data et labels
    data = [item[0] for item in batch]
    labels = [item[1] for item in batch]

    # Custom processing
    # Ex: padding pour séquences de longueurs différentes
    max_len = max(len(d) for d in data)
    padded_data = [torch.nn.functional.pad(d, (0, max_len - len(d))) for d in data]

    # Stack
    data_tensor = torch.stack(padded_data)
    labels_tensor = torch.tensor(labels)

    return data_tensor, labels_tensor

# Utilisation
dataloader = DataLoader(
    dataset,
    batch_size=32,
    collate_fn=custom_collate
)

Datasets pré-construits avec torchvision

PyTorch fournit des datasets populaires prêts à l’emploi.

MNIST et CIFAR

from torchvision import datasets, transforms

# Transformations
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST mean, std
])

# MNIST
mnist_train = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

mnist_test = datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transform
)

# CIFAR-10
cifar10_train = datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

# Créer DataLoaders
train_loader = DataLoader(mnist_train, batch_size=64, shuffle=True, num_workers=2)
test_loader = DataLoader(mnist_test, batch_size=64, shuffle=False, num_workers=2)

# Utilisation
for images, labels in train_loader:
    print(f"Batch: images.shape={images.shape}, labels.shape={labels.shape}")
    break
# Batch: images.shape=torch.Size([64, 1, 28, 28]), labels.shape=torch.Size([64])

ImageNet et ImageFolder

from torchvision import datasets

# ImageFolder : structure train/class_0/, train/class_1/, ...
train_dataset = datasets.ImageFolder(
    root='data/train/',
    transform=train_transform
)

# Accéder aux classes
print(f"Classes: {train_dataset.classes}")  # ['cat', 'dog', ...]
print(f"Class to idx: {train_dataset.class_to_idx}")  # {'cat': 0, 'dog': 1, ...}

# ImageNet (nécessite téléchargement manuel)
imagenet_train = datasets.ImageNet(
    root='/path/to/imagenet/',
    split='train',
    transform=train_transform
)

Gérer les classes déséquilibrées

Problème courant : dataset avec 90% classe A, 10% classe B.

WeightedRandomSampler

from torch.utils.data import WeightedRandomSampler

# Calculer les poids par classe
class_counts = [4500, 500]  # Classe 0: 4500, Classe 1: 500
class_weights = 1. / torch.tensor(class_counts, dtype=torch.float)

# Assigner un poids à chaque sample selon sa classe
sample_weights = torch.zeros(len(dataset))
for idx in range(len(dataset)):
    _, label = dataset[idx]
    sample_weights[idx] = class_weights[label]

# Créer le sampler
sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(sample_weights),
    replacement=True  # Avec remplacement (peut piocher même sample 2×)
)

# DataLoader avec sampler (pas de shuffle!)
dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,  # Utilise sampler au lieu de shuffle
    num_workers=4
)

# Résultat : batches équilibrés ~50% classe 0, ~50% classe 1

Weighted Loss

Alternative : garder le dataset déséquilibré mais pondérer la loss :

import torch.nn as nn

# Calculer poids inverses
class_counts = torch.tensor([4500, 500])
class_weights = 1.0 / class_counts
class_weights = class_weights / class_weights.sum()  # Normaliser

# Loss avec poids
criterion = nn.CrossEntropyLoss(weight=class_weights)

Subset et Split : Train/Val/Test

Diviser un dataset en train/validation/test :

from torch.utils.data import random_split, Subset

# Méthode 1 : random_split
dataset = MyDataset(...)
train_size = int(0.7 * len(dataset))
val_size = int(0.15 * len(dataset))
test_size = len(dataset) - train_size - val_size

train_dataset, val_dataset, test_dataset = random_split(
    dataset,
    [train_size, val_size, test_size],
    generator=torch.Generator().manual_seed(42)  # Reproductible
)

# Méthode 2 : indices manuels
import numpy as np

indices = np.arange(len(dataset))
np.random.shuffle(indices)

train_indices = indices[:train_size]
val_indices = indices[train_size:train_size+val_size]
test_indices = indices[train_size+val_size:]

train_dataset = Subset(dataset, train_indices)
val_dataset = Subset(dataset, val_indices)
test_dataset = Subset(dataset, test_indices)

# Créer DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

Exemple complet : Dataset custom avec augmentation

Mettons tout ensemble :

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os

class CustomImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.images = []
        self.labels = []

        # Charger chemins et labels
        for label, class_name in enumerate(sorted(os.listdir(root_dir))):
            class_path = os.path.join(root_dir, class_name)
            if not os.path.isdir(class_path):
                continue
            for img_name in os.listdir(class_path):
                if img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                    self.images.append(os.path.join(class_path, img_name))
                    self.labels.append(label)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = Image.open(self.images[idx]).convert('RGB')
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

# Transformations
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Datasets
train_dataset = CustomImageDataset('data/train/', transform=train_transform)
val_dataset = CustomImageDataset('data/val/', transform=val_transform)

# DataLoaders
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True
)

val_loader = DataLoader(
    val_dataset,
    batch_size=32,
    shuffle=False,
    num_workers=4,
    pin_memory=True
)

# Utilisation
for epoch in range(10):
    # Training
    for images, labels in train_loader:
        # images: (32, 3, 224, 224)
        # labels: (32,)
        # ... training ...
        pass

    # Validation
    for images, labels in val_loader:
        # ... validation ...
        pass

Exercices pratiques

Dataset custom depuis CSV

Créez un dataset qui charge des images depuis un fichier CSV :

# CSV format :
# image_path,label
# train/img1.jpg,0
# train/img2.jpg,1

class CSVDataset(Dataset):
    def __init__(self, csv_file, root_dir, transform=None):
        # Votre code ici
        pass

    def __len__(self):
        # Votre code ici
        pass

    def __getitem__(self, idx):
        # Votre code ici
        pass

Comparer num_workers

Mesurez l’impact de num_workers sur votre machine :

import time

for num_workers in [0, 2, 4, 8]:
    loader = DataLoader(dataset, batch_size=32, num_workers=num_workers)
    start = time.time()
    for _ in loader:
        pass
    print(f"num_workers={num_workers}: {time.time()-start:.2f}s")

Data Augmentation Visualisation

Visualisez l’effet de la data augmentation :

import matplotlib.pyplot as plt

dataset = CustomImageDataset('data/', transform=train_transform)

# Charger même image plusieurs fois pour voir variations
img, label = dataset[0]

fig, axes = plt.subplots(2, 4, figsize=(12, 6))
for i in range(8):
    img_aug, _ = dataset[0]  # Même index mais augmentation différente
    ax = axes[i//4, i%4]
    ax.imshow(img_aug.permute(1, 2, 0))  # (C, H, W) → (H, W, C)
    ax.axis('off')
plt.show()

Conclusion

Dans cet article, vous avez maîtrisé la gestion des données avec PyTorch :

✅ Dataset custom (hériter de Dataset, len, getitem)
✅ DataLoader (batching, shuffling, num_workers)
✅ Transformations avec torchvision.transforms
✅ Data augmentation (RandomCrop, Flip, ColorJitter)
✅ Différences train vs validation transforms
✅ Datasets pré-construits (MNIST, CIFAR, ImageNet)
✅ Gérer classes déséquilibrées (WeightedRandomSampler)
✅ Optimisations (pin_memory, persistent_workers)
✅ Split train/val/test

Prochaine étape : Dans le prochain article, nous allons découvrir le transfer learning ! Vous apprendrez à utiliser des modèles pré-entraînés (ResNet, EfficientNet, Vision Transformers) et à les fine-tuner sur vos propres données.

Pour aller plus loin :

Explorez la documentation Dataset/DataLoader
Découvrez torchvision.transforms en détail
Consultez notre guide sur les CNNs pour computer vision avancée

🔎 Tip
Exercice complet : Créez un dataset custom pour un problème de classification d’images. Implémentez data augmentation, split train/val, et optimisez le DataLoader. Mesurez l’impact de num_workers et pin_memory sur les performances. Visualisez les augmentations appliquées.