Learn AI with Python · Lesson

Custom Datasets and DataLoaders

torch.utils.data.Dataset, __len__/__getitem__, DataLoader, transforms, augmentation.

Custom Datasets and DataLoaders is a free Learn AI with Python lesson on CoddyKit — lesson 2 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the Learn AI with Python learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

Feeding Data to a Model

Training needs an efficient way to read, transform, and batch data. PyTorch provides two abstractions: Dataset (knows how to fetch one sample) and DataLoader (batches and shuffles them).

from torch.utils.data import Dataset, DataLoader

The Dataset Interface

A custom Dataset subclass must implement two methods: __len__ (how many samples) and __getitem__ (return the sample at an index). PyTorch calls these to pull data.

Implementing len

__len__ tells PyTorch the dataset size so it knows how many indices exist and how many batches an epoch contains.

class ImageDataset(Dataset):
    def __init__(self, paths, labels):
        self.paths = paths
        self.labels = labels

    def __len__(self):
        return len(self.paths)

Implementing getitem

__getitem__ loads and returns one sample (and its label) given an index. This is where you open an image file and convert it into a tensor.

from PIL import Image

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert("RGB")
        label = self.labels[idx]
        return img, label

Why Transforms?

Raw images vary in size and pixel range. Transforms standardize them: resize to a fixed shape, convert to a tensor, and normalize pixel values so the model trains stably.

from torchvision import transforms

transforms.Compose

transforms.Compose chains several transforms into one pipeline applied in order. A typical chain is Resize then ToTensor then Normalize.

tf = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

Understanding the Transforms

Resize fixes the spatial size; ToTensor converts a PIL image to a tensor and scales pixels to [0,1]; Normalize shifts and scales each channel to zero mean and unit variance, which speeds convergence.

Applying Transforms in the Dataset

Pass the transform into the dataset and apply it inside __getitem__ so every sample is preprocessed consistently as it is fetched.

class ImageDataset(Dataset):
    def __init__(self, paths, labels, transform):
        self.paths, self.labels, self.transform = paths, labels, transform

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert("RGB")
        return self.transform(img), self.labels[idx]

Wrapping in a DataLoader

The DataLoader turns a Dataset into an iterable of batches. Set batch_size to control samples per step and shuffle=True to randomize order each epoch (important for training).

dataset = ImageDataset(paths, labels, tf)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

num_workers for Speed

num_workers spawns parallel subprocesses to load and transform data while the GPU trains, hiding I/O latency. A value like 4 often keeps the GPU fed instead of waiting.

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

Iterating Batches

Loop over the DataLoader to get batched tensors. Each iteration yields (images, labels) where images has shape [batch_size, channels, H, W], ready for the model.

for images, labels in loader:
    print(images.shape)  # torch.Size([32, 3, 224, 224])
    break

Quick Check

Test your data pipeline knowledge.

Recap: Datasets and DataLoaders

You built a custom Dataset with __len__ and __getitem__, preprocessed images using transforms.Compose (Resize, ToTensor, Normalize), and wrapped it in a DataLoader with batch_size, shuffle, and num_workers to feed batches efficiently to your model.

Frequently asked questions

Is the “Custom Datasets and DataLoaders” lesson free?

Yes — the full text of “Custom Datasets and DataLoaders” is free to read here on the web, and the Learn AI with Python course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the Learn AI with Python course, upgrade to CoddyKit PRO.

What will I learn in “Custom Datasets and DataLoaders”?

torch.utils.data.Dataset, __len__/__getitem__, DataLoader, transforms, augmentation. You practise Learn AI with Python with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start Learn AI with Python?

No prior experience is required. Learn AI with Python on CoddyKit is structured for beginners through advanced learners; this is — lesson 2 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Custom Datasets and DataLoaders” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this Learn AI with Python lesson?

Yes. Every Learn AI with Python lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

← Back to Learn AI with Python