How to randomly split a dataset in PyTorch

Training and validation data need a fixed boundary before model metrics mean anything. PyTorch can build that boundary with random_split(), which chooses non-overlapping source indices from a map-style dataset and returns subset objects that still read from the original samples.

A seeded torch.Generator makes the random split repeatable. Reusing the same seed keeps the same samples in the training and validation subsets across runs, while changing the seed intentionally reshuffles the membership.

Use exact integer lengths when a project needs fixed sample counts. Fraction lengths are accepted by random_split() when they sum to 1, but integer lengths make a smoke test easier to review because the requested counts and printed subset sizes match directly.

1. Create a smoke script that builds a map-style dataset, splits it, and checks the subset membership.

   split_dataset.py

   import torch
   from torch.utils.data import DataLoader, TensorDataset, random_split
    
   features = torch.arange(0, 60, dtype=torch.float32).reshape(20, 3)
   labels = torch.arange(20)
   dataset = TensorDataset(features, labels)
    
   generator = torch.Generator().manual_seed(20260707)
   train_set, valid_set = random_split(dataset, [14, 6], generator=generator)
    
   repeat_train, repeat_valid = random_split(
       dataset,
       [14, 6],
       generator=torch.Generator().manual_seed(20260707),
   )
    
   train_indices = list(train_set.indices)
   valid_indices = list(valid_set.indices)
   repeat_train_indices = list(repeat_train.indices)
    
   print(f"dataset length: {len(dataset)}")
   print(f"train length: {len(train_set)}")
   print(f"validation length: {len(valid_set)}")
   print(f"train indices: {train_indices}")
   print(f"validation indices: {valid_indices}")
   print(f"overlap count: {len(set(train_indices) & set(valid_indices))}")
   print(f"repeat matches: {train_indices == repeat_train_indices}")
    
   train_loader = DataLoader(train_set, batch_size=4, shuffle=True)
   valid_loader = DataLoader(valid_set, batch_size=6, shuffle=False)
   train_batch, _ = next(iter(train_loader))
   _, valid_labels = next(iter(valid_loader))
    
   print(f"train batch shape: {tuple(train_batch.shape)}")
   print(f"validation labels: {valid_labels.tolist()}")

   random_split() returns Subset objects. The indices attribute is useful for a smoke test, but training code normally passes the subsets directly to DataLoader.

2. Run the split script in the Python environment that has PyTorch installed.

   $ python3 split_dataset.py
   dataset length: 20
   train length: 14
   validation length: 6
   train indices: [9, 0, 19, 10, 18, 16, 5, 15, 8, 17, 3, 11, 12, 6]
   validation indices: [7, 2, 14, 13, 4, 1]
   overlap count: 0
   repeat matches: True
   train batch shape: (4, 3)
   validation labels: [7, 2, 14, 13, 4, 1]

3. Confirm that the requested split lengths match the printed subset lengths.

   The smoke script uses 14 and 6 from a 20-sample dataset. For a real dataset, make the values add up to len(dataset) before training.

4. Confirm that overlap count is 0 and repeat matches is True.

   A zero overlap means no source index landed in both subsets. A true repeat check means the seeded generator produced the same training subset when the split was recreated.

5. Set loader shuffling only where the phase needs it.

   The training DataLoader can use shuffle=True for batch order. The validation DataLoader should usually keep shuffle=False so evaluation output follows a stable order.

   Related: How to use a DataLoader in PyTorch

6. Remove the smoke script when the split pattern has been copied into project code.

   $ rm split_dataset.py