Introduction

torchvision.datasets.ImageFolder loads image samples from a directory tree, but it does not split the dataset into training and test sets for you. The usual workflow is to build the dataset once, create a reproducible index split, and then expose those indices through Subset objects. The subtle part is transforms: train and test sets often need different transforms, so sharing one dataset object is not always enough.

Start with One ImageFolder

An ImageFolder expects this directory layout:

dataset/
  cats/
  dogs/

Each subdirectory name becomes a class label. To inspect the full dataset:

1from torchvision import datasets, transforms
2
3base_transform = transforms.Compose([
4    transforms.Resize((224, 224)),
5    transforms.ToTensor(),
6])
7
8full_dataset = datasets.ImageFolder("dataset", transform=base_transform)
9print(len(full_dataset))
10print(full_dataset.classes)

At this point, you have one dataset containing all samples.

Use random_split for a Simple Split

If train and test can share the same transform, random_split is the simplest solution.

1import torch
2from torch.utils.data import random_split
3
4generator = torch.Generator().manual_seed(42)
5
6train_size = int(0.8 * len(full_dataset))
7test_size = len(full_dataset) - train_size
8
9train_dataset, test_dataset = random_split(
10    full_dataset,
11    [train_size, test_size],
12    generator=generator,
13)
14
15print(len(train_dataset), len(test_dataset))

This works well for many baseline projects and gives a reproducible split because the generator seed is fixed.

The Transform Problem

In real training pipelines, train and test usually need different transforms. Training often uses augmentation, while test uses only deterministic preprocessing.

That creates a problem: random_split returns subsets that still point at the same underlying dataset object. If that dataset has one transform, both subsets use it.

For example, this is often not what you want:

• training with random flips and crops
• test set also receiving random flips and crops accidentally

The fix is to separate the indices from the dataset object.

Use Two Dataset Instances with Shared Indices

Create two ImageFolder datasets that point at the same directory but use different transforms. Then apply the same split indices to both.

1import torch
2from torch.utils.data import Subset
3from torchvision import datasets, transforms
4
5train_transform = transforms.Compose([
6    transforms.Resize((224, 224)),
7    transforms.RandomHorizontalFlip(),
8    transforms.ToTensor(),
9])
10
11test_transform = transforms.Compose([
12    transforms.Resize((224, 224)),
13    transforms.ToTensor(),
14])
15
16train_source = datasets.ImageFolder("dataset", transform=train_transform)
17test_source = datasets.ImageFolder("dataset", transform=test_transform)
18
19generator = torch.Generator().manual_seed(42)
20indices = torch.randperm(len(train_source), generator=generator).tolist()
21
22split = int(0.8 * len(indices))
23train_indices = indices[:split]
24test_indices = indices[split:]
25
26train_dataset = Subset(train_source, train_indices)
27test_dataset = Subset(test_source, test_indices)

This gives reproducible sample membership while keeping train and test preprocessing separate.

Build DataLoaders After the Split

Once the subsets exist, wrap them in DataLoader objects.

1from torch.utils.data import DataLoader
2
3train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
4test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

Shuffle the training loader, not the test loader. That keeps evaluation deterministic and easier to debug.

Think About Class Balance

random_split does not guarantee stratification. If the dataset is imbalanced, the train and test subsets may end up with noticeably different class ratios. For many projects that is acceptable. For more sensitive evaluation, stratified splitting is better, but that usually requires building the split from the label list explicitly.

With ImageFolder, labels are available in dataset.targets, so class-aware splitting is possible if needed. The important point is that random splitting and balanced splitting are not the same thing.

Common Pitfalls

• Using random_split on one dataset object and then forgetting that both subsets share the same transform.
• Applying training augmentation to the test set by accident.
• Forgetting to fix the random seed and then getting different splits on every run.
• Shuffling the test loader and making evaluation harder to compare.
• Assuming a random split is automatically stratified by class.

Summary

• 'ImageFolder loads the full dataset, but you still need to define the train-test split yourself.'
• 'random_split is fine when both subsets can share the same transform.'
• When train and test need different transforms, use separate dataset instances with shared split indices.
• Build DataLoader objects after creating the subsets.
• If class balance matters, do not assume a random split is enough.