Load csv and Image dataset in pytorch

pytorch

csv

image dataset

data loading

machine learning

Load csv and Image dataset in pytorch

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

PyTorch Machine Learning

Introduction

PyTorch loads data through the Dataset and DataLoader classes. For CSV data, you create a custom Dataset that reads a CSV file (typically with pandas) and returns tensors. For image datasets, use torchvision.datasets.ImageFolder for directory-structured data or a custom Dataset that reads image paths from a CSV and loads them with PIL. The DataLoader then wraps any Dataset to provide batching, shuffling, and parallel data loading. This two-layer design separates data access logic from training loop mechanics.

Loading a CSV Dataset

python

1import torch
2from torch.utils.data import Dataset, DataLoader
3import pandas as pd
4
5class CSVDataset(Dataset):
6    def __init__(self, csv_file, target_column):
7        self.data = pd.read_csv(csv_file)
8        self.features = self.data.drop(columns=[target_column]).values
9        self.targets = self.data[target_column].values
10
11    def __len__(self):
12        return len(self.data)
13
14    def __getitem__(self, idx):
15        x = torch.tensor(self.features[idx], dtype=torch.float32)
16        y = torch.tensor(self.targets[idx], dtype=torch.float32)
17        return x, y
18
19# Usage
20dataset = CSVDataset("train.csv", target_column="price")
21loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=2)
22
23for batch_x, batch_y in loader:
24    print(batch_x.shape)  # torch.Size([32, num_features])
25    print(batch_y.shape)  # torch.Size([32])
26    break

The __getitem__ method returns a single sample as tensors. DataLoader collates individual samples into batches automatically.

Loading Images from a Directory

python

1from torchvision import datasets, transforms
2
3# Directory structure:
4# data/train/
5#   cats/
6#     cat001.jpg
7#     cat002.jpg
8#   dogs/
9#     dog001.jpg
10#     dog002.jpg
11
12transform = transforms.Compose([
13    transforms.Resize((224, 224)),
14    transforms.ToTensor(),
15    transforms.Normalize(mean=[0.485, 0.456, 0.406],
16                         std=[0.229, 0.224, 0.225])
17])
18
19dataset = datasets.ImageFolder(root="data/train", transform=transform)
20loader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=4)
21
22print(dataset.classes)       # ['cats', 'dogs']
23print(dataset.class_to_idx)  # {'cats': 0, 'dogs': 1}
24
25for images, labels in loader:
26    print(images.shape)  # torch.Size([16, 3, 224, 224])
27    print(labels.shape)  # torch.Size([16])
28    break

ImageFolder automatically assigns class labels based on subdirectory names. The transform pipeline handles resizing, tensor conversion, and normalization.

Loading Images Referenced by a CSV

python

1import torch
2from torch.utils.data import Dataset, DataLoader
3from PIL import Image
4from torchvision import transforms
5import pandas as pd
6import os
7
8class CSVImageDataset(Dataset):
9    """CSV contains image paths and labels."""
10
11    def __init__(self, csv_file, img_dir, transform=None):
12        self.annotations = pd.read_csv(csv_file)
13        self.img_dir = img_dir
14        self.transform = transform
15
16    def __len__(self):
17        return len(self.annotations)
18
19    def __getitem__(self, idx):
20        img_path = os.path.join(
21            self.img_dir,
22            self.annotations.iloc[idx, 0]  # First column: filename
23        )
24        image = Image.open(img_path).convert("RGB")
25        label = torch.tensor(self.annotations.iloc[idx, 1])  # Second column: label
26
27        if self.transform:
28            image = self.transform(image)
29
30        return image, label
31
32# CSV format:
33# filename,label
34# img001.jpg,0
35# img002.jpg,1
36
37transform = transforms.Compose([
38    transforms.Resize((224, 224)),
39    transforms.RandomHorizontalFlip(),
40    transforms.ToTensor(),
41])
42
43dataset = CSVImageDataset(
44    csv_file="labels.csv",
45    img_dir="data/images/",
46    transform=transform
47)
48loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

This pattern is common when labels are in a CSV file separate from the image directory, or when you need metadata beyond simple class folders.

Combined CSV Features + Images

python

1class MultiModalDataset(Dataset):
2    """Dataset that combines tabular features from CSV with image data."""
3
4    def __init__(self, csv_file, img_dir, transform=None):
5        self.df = pd.read_csv(csv_file)
6        self.img_dir = img_dir
7        self.transform = transform
8
9    def __len__(self):
10        return len(self.df)
11
12    def __getitem__(self, idx):
13        row = self.df.iloc[idx]
14
15        # Tabular features
16        features = torch.tensor(
17            [row["age"], row["price"], row["rating"]],
18            dtype=torch.float32
19        )
20
21        # Image
22        img_path = os.path.join(self.img_dir, row["image_file"])
23        image = Image.open(img_path).convert("RGB")
24        if self.transform:
25            image = self.transform(image)
26
27        # Label
28        label = torch.tensor(row["label"], dtype=torch.long)
29
30        return features, image, label

For multi-modal models, a single Dataset can return multiple inputs. The DataLoader handles batching all of them together.

DataLoader Configuration

python

1loader = DataLoader(
2    dataset,
3    batch_size=64,
4    shuffle=True,           # Shuffle for training, False for validation
5    num_workers=4,          # Parallel data loading processes
6    pin_memory=True,        # Faster CPU-to-GPU transfer
7    drop_last=True,         # Drop incomplete final batch
8    persistent_workers=True # Keep workers alive between epochs (PyTorch 1.8+)
9)
10
11# Iterate
12for epoch in range(num_epochs):
13    for batch in loader:
14        features, labels = batch
15        features = features.to("cuda")
16        labels = labels.to("cuda")
17        # Training step...

Common Pitfalls

Forgetting __len__ or __getitem__: DataLoader requires both methods on the dataset. Missing either raises TypeError. __len__ returns the dataset size, __getitem__ returns a single sample by index.
Not converting to tensors in __getitem__: Returning raw NumPy arrays or Python lists works but is slower. PyTorch's default collate function converts them to tensors, but explicit conversion in __getitem__ catches type errors earlier and is more efficient.
Using num_workers > 0 on Windows without if __name__ == '__main__': Windows uses spawn for multiprocessing, which re-imports the module. Without the __main__ guard, workers crash with pickle errors. Always wrap DataLoader iteration in if __name__ == '__main__' on Windows.
Loading all images into memory at __init__: For large image datasets, loading all images during initialization causes out-of-memory errors. Load images lazily in __getitem__ — each image is loaded on demand and garbage-collected after use.
Not normalizing image tensors: ToTensor() scales pixel values to [0, 1], but pretrained models (ResNet, VGG) expect ImageNet normalization (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]). Missing this step causes poor model performance with pretrained weights.

Summary

Create a custom Dataset subclass with __len__ and __getitem__ for CSV data
Use torchvision.datasets.ImageFolder for directory-structured image datasets
For CSV-referenced images, build a custom Dataset that reads paths from the CSV and loads images in __getitem__
Wrap any Dataset with DataLoader for batching, shuffling, and parallel loading
Use pin_memory=True and num_workers > 0 for faster GPU training
Apply transforms.Compose for image preprocessing (resize, normalize, augment)

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Load csv and Image dataset in pytorch

ML System Design practice on Codemia

Introduction

Loading a CSV Dataset

Loading Images from a Directory

Loading Images Referenced by a CSV

Combined CSV Features + Images

DataLoader Configuration

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia