PyTorch DataLoader : lots, mélange et processus de chargement

Un DataLoader PyTorch fournit des lots d’exemples à une boucle d’entraînement ou d’évaluation. Le jeu de données définit les exemples ; le chargeur organise leur regroupement et leur livraison. Avant d’optimiser sa vitesse, vérifiez la forme des tenseurs, leur type et la présence des exemples attendus.

Ce guide part d’un petit exemple sur CPU, sans téléchargement de données. Il couvre les lots, le mélange, les processus de chargement et les entrées de longueurs différentes, avec des contrôles qui rendent visibles les exemples manquants ou répétés.

Dataset et DataLoader : deux responsabilités

Un Dataset indexable définit __len__ et __getitem__ : le nombre d’exemples et la façon d’en récupérer un par son indice. Un DataLoader regroupe ces exemples en lots. Le tutoriel officiel sur les jeux de données présente cette séparation.

Chaque exemple ci-dessous est un dictionnaire contenant trois caractéristiques, une étiquette de classe et un identifiant. Cet identifiant sert au diagnostic. Le conserver pendant un petit essai permet de vérifier la couverture des données indépendamment de la perte du modèle.

Utilisez un environnement où PyTorch est installé et notez sa version avant de comparer les résultats. Notre guide de l’environnement de développement explique comment isoler les dépendances d’un projet.

Un exemple DataLoader PyTorch exécutable

Enregistrez ce code dans dataloader_demo.py. Il crée neuf exemples synthétiques et demande des lots de quatre. Aucun GPU, modèle ou fichier de données externe n’est nécessaire.

import argparse
import torch
from torch.utils.data import DataLoader, Dataset


class ToyDataset(Dataset):
    def __len__(self):
        return 9

    def __getitem__(self, index):
        return {
            'features': torch.tensor([index, index + 1, index + 2],
                                     dtype=torch.float32),
            'label': torch.tensor(index % 2, dtype=torch.long),
            'sample_id': index,
        }


def collect(workers=0, shuffle=True, drop_last=False, seed=11):
    loader = DataLoader(
        ToyDataset(), batch_size=4, shuffle=shuffle,
        num_workers=workers, drop_last=drop_last,
        generator=torch.Generator().manual_seed(seed),
    )
    sizes, ids = [], []
    for batch in loader:
        x, y = batch['features'], batch['label']
        assert x.shape == (len(y), 3)
        assert x.dtype == torch.float32 and y.dtype == torch.long
        sizes.append(len(y))
        ids.extend(batch['sample_id'].tolist())
    return sizes, ids


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--workers', type=int, default=0)
    args = parser.parse_args()
    sizes, ids = collect(workers=args.workers)
    assert sorted(ids) == list(range(9))
    print('batch sizes:', sizes)
    print('samples:', len(ids), 'unique:', len(set(ids)))

Exécutez le fichier avec un chargement dans le processus principal, puis avec deux processus de chargement :

python dataloader_demo.py --workers 0
python dataloader_demo.py --workers 2

Les deux commandes ont produit ce résultat lors de notre test :

batch sizes: [4, 4, 1]
samples: 9 unique: 9

Les deux premiers lots contiennent quatre exemples chacun, et le dernier en contient un. Les neuf identifiants apparaissent exactement une fois sur ces trois lots. L’assertion de forme vérifie que le tenseur des caractéristiques a trois colonnes et autant de lignes qu’il y a d’étiquettes. Le code vérifie aussi le type flottant des caractéristiques et le type entier des étiquettes.

Ces choix conviennent à ce petit jeu de classification. Votre tâche peut demander des images, des identifiants de tokens, des cibles de régression ou des masques supplémentaires. Examinez un lot réel avant d’ajouter le modèle, plutôt que de reprendre ces dimensions dans un autre entraînement.

Comprendre la taille des lots et le dernier lot

Nous avons testé le même chargeur avec drop_last=True. Il a produit deux lots de quatre et huit identifiants distincts. Sans mélange, l’identifiant 8 était omis ; avec la graine et le mélange utilisés ici, c’était l’identifiant 3. Ce résultat invite à vérifier les exemples eux-mêmes, et pas seulement le nombre de lots.

Conservez le dernier lot partiel lorsque tous les exemples doivent contribuer. Si votre entraînement exige des lots de taille égale, rendez l’omission explicite et consignez-la. Pour l’évaluation, vérifiez la couverture complète et utilisez les bons dénominateurs pour les mesures agrégées. Faire une moyenne non pondérée entre la moyenne d’un lot de quatre exemples et celle d’un lot d’un exemple donnerait trop de poids au dernier exemple.

La taille des lots modifie aussi le nombre d’itérations. Dans ce test de neuf exemples, des lots de quatre donnent trois itérations si le dernier lot est conservé. Revoyez le fonctionnement du programme de taux d’apprentissage, des journaux ou des intervalles de checkpoint lorsqu’ils comptent les itérations.

Mélanger les données sans mélanger les sous-ensembles

L’exemple active shuffle=True avec un générateur initialisé par une graine. Nos tests ont confirmé que les neuf identifiants restaient présents avec drop_last=False. Le mélange changeait leur ordre, pas leur appartenance au jeu de données.

Définissez séparément les ensembles d’entraînement, de validation et de test avant de construire les chargeurs. Un chargeur qui mélange les données ne crée pas une séparation valide. Pour des séries temporelles, des observations groupées ou plusieurs enregistrements d’une personne, respectez ces relations lors de la séparation. Conservez des identifiants par sous-ensemble pour détecter les chevauchements accidentels.

Un nouveau chargeur avec un nouveau générateur initialisé à 11 a reproduit notre ordre initial. Réutiliser un générateur fait avancer son état. Les transformations aléatoires peuvent demander d’autres graines, notamment pour Python ou NumPy via worker_init_fn. Suivez le guide de reproductibilité PyTorch : une graine identique ne garantit pas les mêmes résultats entre versions ou plateformes.

Choisir num_workers après avoir validé le chargement

num_workers=0 charge les données dans le processus principal. Une valeur positive utilise des sous-processus. Commencez à zéro pour déboguer, puis comparez quelques processus sur des entrées représentatives. Le guide d’optimisation PyTorch décrit ce chargement et la mémoire épinglée.

Notre exemple synthétique effectue très peu de préparation. Sa réussite avec deux processus établit que le code fonctionne dans cette configuration, sans démontrer un gain de vitesse. Un essai réel doit inclure vos lectures de fichiers, le décodage, les transformations et la boucle d’entraînement.

Notez la durée, le nombre d’exemples traités, l’utilisation du CPU et celle de la mémoire hôte pour une même charge. Répétez la mesure et distinguez le démarrage du fonctionnement établi. Si davantage de processus augmente la mémoire sans réduire la durée totale, conservez la configuration la plus simple. Un GPU rapide ne corrige pas à lui seul une source de données lente.

Exécuter les exemples avec plusieurs processus sous Windows

Windows démarre les processus avec spawn. Définissez la classe du jeu de données et la fonction de regroupement au niveau du module, puis lancez l’exécution sous if __name__ == '__main__':. La FAQ Windows de PyTorch explique cette protection et conseille de revenir à zéro processus de chargement pour diagnostiquer leurs erreurs.

L’exemple respecte cette structure. S’il fonctionne dans un fichier enregistré mais échoue dans un notebook, reproduisez d’abord le problème dans le fichier autonome avec le même interpréteur. Conservez la trace d’erreur complète et réduisez les données au plus petit cas qui échoue.

Utiliser collate_fn pour des entrées de longueurs différentes

Le premier jeu de données fournit des vecteurs de taille fixe. Pour des séquences de longueurs différentes, définissez comment construire un lot rectangulaire. Ce second exemple complète deux séquences et renvoie leurs longueurs initiales avec leurs étiquettes. Enregistrez-le dans padding_demo.py et exécutez-le avec Python.

import torch
from torch.nn.utils.rnn import pad_sequence


def pad_batch(samples):
    sequences, labels = zip(*samples)
    lengths = torch.tensor([len(x) for x in sequences])
    padded = pad_sequence(sequences, batch_first=True, padding_value=0)
    return padded, torch.tensor(labels, dtype=torch.long), lengths


if __name__ == '__main__':
    samples = [(torch.tensor([4, 8, 2]), 1), (torch.tensor([5]), 0)]
    padded, labels, lengths = pad_batch(samples)
    print(padded.tolist(), labels.tolist(), lengths.tolist())

Les valeurs complétées sont [[4, 8, 2], [5, 0, 0]], les étiquettes sont [1, 0] et les longueurs sont [3, 1]. La référence pad_sequence documente la disposition batch-first et la valeur de remplissage.

Fournissez collate_fn=pad_batch au chargeur d’un jeu de données qui renvoie ces paires séquence-étiquette. Notre vérification a aussi utilisé cette fonction dans un DataLoader à deux processus. Gardez-la au niveau du module pour que les processus puissent l’importer.

Le remplissage ajoute des positions qui ne sont pas des observations d’origine. Utilisez les longueurs renvoyées pour construire le masque ou la représentation attendue par le modèle. Excluez le remplissage du calcul de la perte lorsque nécessaire et choisissez une valeur adaptée à votre vocabulaire. Le zéro numérique est un choix d’exemple, pas une règle universelle de token de remplissage.

Séparer chargement CPU et transfert GPU

pin_memory=True prépare des tampons adaptés en mémoire hôte ; il ne transfère pas le lot sur un GPU. Ce transfert reste une étape distincte. Pour CUDA, évaluez ce réglage avec des transferts utilisant non_blocking=True en suivant le tutoriel PyTorch sur les transferts. Le bénéfice dépend de la charge, et la réutilisation des tampons demande de l’attention avec les transferts asynchrones.

Validez l’exemple CPU avant d’ajouter le placement sur un appareil. Notre guide PyTorch sur A100 couvre l’installation CUDA et les vérifications du GPU. Nous n’avons testé ni les transferts GPU ni son taux d’utilisation pour cet article.

Régler les options avancées une à une

prefetch_factor compte les lots préparés à l’avance par processus. persistent_workers=True conserve les processus et leurs instances du jeu de données entre les passages. Ces réglages peuvent modifier la mémoire hôte et la durée de vie de l’état. Consultez l’API DataLoader de votre version avant de les changer.

Ce guide utilise un jeu indexable dans un seul processus d’entraînement. Les flux IterableDataset et l’échantillonnage distribué exigent leurs propres contrôles de répartition et de couverture. Ajouter des processus de chargement ou des rangs d’entraînement ne garantit pas de conserver un seul exemplaire de chaque identifiant.

Vérifications avant une longue exécution

• Examinez les formes, les types, les étiquettes et le masque de remplissage éventuel.
• Comptez les identifiants sur un passage complet et expliquez les omissions ou répétitions.
• Vérifiez la composition des sous-ensembles séparément de l’ordre des exemples.
• Reproduisez les erreurs avec zéro processus de chargement et peu d’entrées.
• Modifiez un réglage de performance à la fois et mesurez la charge complète.

Nos dix contrôles CPU ont réussi sous Windows avec Python 3.12.14 et PyTorch 2.15.0+cpu. Ils couvraient zéro et deux processus de chargement, le mélange activé ou non, la conservation ou l’abandon du dernier lot, la répétition de l’ordre avec une même graine et le remplissage personnalisé avec deux processus. Ces contrôles portent sur la correction du traitement de données synthétiques. Ils ne mesurent ni les entrées-sorties de stockage, ni CUDA, ni la vitesse d’entraînement.