Introduction

To train a CNN in PyTorch, you need a dataset object that can return one image-label pair at a time and a DataLoader that can batch those samples efficiently. The best implementation depends on how your files are organized: a simple folder-per-class layout can use ImageFolder, while CSV-driven or irregular datasets usually need a custom Dataset.

The important part is not only making the code run. You also need the transforms, labels, and tensor shapes to match the network you plan to train.

Use ImageFolder When the Layout Fits

If your dataset is organized by class directory, torchvision.datasets.ImageFolder already does most of the work:

1dataset/
2  train/
3    cats/
4    dogs/
5  val/
6    cats/
7    dogs/

Then you can load it like this:

1from torchvision import datasets, transforms
2from torch.utils.data import DataLoader
3
4train_transforms = transforms.Compose([
5    transforms.Resize((224, 224)),
6    transforms.ToTensor(),
7])
8
9train_dataset = datasets.ImageFolder("dataset/train", transform=train_transforms)
10val_dataset = datasets.ImageFolder("dataset/val", transform=train_transforms)
11
12train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
13val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
14
15print(train_dataset.class_to_idx)

If you control the dataset layout, this is often the simplest answer. It saves you from writing label-discovery code manually.

Write a Custom Dataset for CSV or Metadata-Based Labels

If image filenames and labels live in a CSV file, define your own Dataset:

1from pathlib import Path
2import pandas as pd
3from PIL import Image
4from torch.utils.data import Dataset
5
6class CustomImageDataset(Dataset):
7    def __init__(self, csv_file, root_dir, transform=None):
8        self.df = pd.read_csv(csv_file)
9        self.root_dir = Path(root_dir)
10        self.transform = transform
11
12    def __len__(self):
13        return len(self.df)
14
15    def __getitem__(self, idx):
16        row = self.df.iloc[idx]
17        image_path = self.root_dir / row["filename"]
18        image = Image.open(image_path).convert("RGB")
19        label = int(row["label"])
20
21        if self.transform is not None:
22            image = self.transform(image)
23
24        return image, label

This gives you full control over file paths, label mapping, filtering, and preprocessing.

Wrap It in a DataLoader

The dataset defines how to read one sample. The DataLoader handles batching, shuffling, and parallel loading:

1from torchvision import transforms
2from torch.utils.data import DataLoader
3
4transform = transforms.Compose([
5    transforms.Resize((224, 224)),
6    transforms.RandomHorizontalFlip(),
7    transforms.ToTensor(),
8])
9
10dataset = CustomImageDataset(
11    csv_file="labels.csv",
12    root_dir="images",
13    transform=transform,
14)
15
16loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)
17
18images, labels = next(iter(loader))
19print(images.shape)
20print(labels.shape)

That one-batch sanity check is worth doing before a long training run. It proves that the CNN will receive tensors with the expected shape and that labels are coming through correctly.

Choose Transforms Deliberately

CNN training depends heavily on preprocessing. At minimum, decide:

• target image size
• RGB versus grayscale handling
• normalization values
• which augmentations belong only in training

If you are using a pretrained model, its expected normalization matters. Training may still run without it, but transfer learning quality can drop noticeably.

Debug the Dataset Before Training

Most failed CNN runs are caused by dataset issues rather than architecture issues. Before training, inspect:

• 'class_to_idx or your custom label mapping'
• one or two actual images after transform
• tensor shape and dtype
• class balance

This catches common mistakes such as corrupted images, wrong labels, grayscale images mixed into RGB pipelines, or augmentation accidentally applied to validation data.

Common Pitfalls

• Writing a custom dataset when ImageFolder already matches the directory structure.
• Returning PIL images instead of tensors because transforms were omitted.
• Using labels in a format that does not match the chosen loss function.
• Applying training augmentations to validation or test data.
• Starting training before checking one batch of images and labels.

Summary

• Use ImageFolder when the dataset is already organized by class folders.
• Write a custom Dataset when labels come from CSV files or custom metadata.
• Use a DataLoader for batching, shuffling, and background workers.
• Treat transforms and normalization as part of the model pipeline.
• Inspect a sample batch before training to catch data issues early.