How to use a DataLoader in PyTorch

Model code usually consumes groups of samples even though a PyTorch Dataset returns one sample at a time. A DataLoader supplies the iterable boundary between those two layers by fetching dataset items, collating compatible fields, and yielding mini-batches.

For a map-style dataset such as TensorDataset, the default collate function stacks corresponding fields along a new first dimension. A dataset that returns feature-label pairs therefore produces one feature tensor and one label tensor during each loader iteration.

A seeded torch.Generator makes shuffled sample order reproducible, while num_workers=0 keeps data loading in the main process. Three iterations should cover the six-row sample dataset exactly once and return feature batches whose first dimension equals the requested batch size.

1. Create the dataset section in dataloader_batch_demo.py.

   dataloader_batch_demo.py

   import torch
   from torch.utils.data import DataLoader, TensorDataset
    
    
   features = torch.arange(24, dtype=torch.float32).reshape(6, 4)
   labels = torch.arange(6)
   dataset = TensorDataset(features, labels)

2. Append a deterministic DataLoader configuration below the dataset.

   loader = DataLoader(
       dataset,
       batch_size=2,
       shuffle=True,
       num_workers=0,
       generator=torch.Generator().manual_seed(17),
   )

   batch_size=2 groups two samples per iteration. shuffle=True randomizes the map-style dataset indexes, and the seeded Generator makes that order repeatable. num_workers=0 fetches samples in the main process.

3. Add the batch iteration and coverage output below the loader.

   loaded_labels = []
    
   for batch_number, (batch_features, batch_labels) in enumerate(loader, start=1):
       loaded_labels.extend(batch_labels.tolist())
       print(
           f"batch {batch_number}: "
           f"features={tuple(batch_features.shape)} "
           f"labels={batch_labels.tolist()}"
       )
    
   print(f"loaded samples: {len(loaded_labels)}")
   print(f"sorted labels: {sorted(loaded_labels)}")

4. Run dataloader_batch_demo.py to verify every sample appears in shuffled mini-batches.

   $ python3 dataloader_batch_demo.py
   batch 1: features=(2, 4) labels=[5, 3]
   batch 2: features=(2, 4) labels=[4, 2]
   batch 3: features=(2, 4) labels=[1, 0]
   loaded samples: 6
   sorted labels: [0, 1, 2, 3, 4, 5]