Custom Dataset class in PyTorch

Introduction

We use a Custom Dataset class to organize and prepare data for machine learning models in PyTorch. It helps load data easily and cleanly.

When you have data in files like images or text and want to load them for training.

When your data needs special processing before feeding into a model.

When you want to use PyTorch's DataLoader to handle batches and shuffling.

When your data is not in a standard format supported by PyTorch datasets.

Syntax

PyTorch

from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

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

    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]
        return sample, label

The __init__ method sets up your data.

The __len__ method returns how many samples you have.

The __getitem__ method gets one sample and its label by index.

Examples

A simple dataset with only data, no labels.

PyTorch

class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

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

    def __getitem__(self, idx):
        return self.data[idx]

Dataset for images with optional transformations.

PyTorch

from torchvision.io import read_image

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

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

    def __getitem__(self, idx):
        image = read_image(self.image_paths[idx])
        if self.transform:
            image = self.transform(image)
        label = self.labels[idx]
        return image, label

Sample Model

This program creates a simple dataset with features and labels, then uses DataLoader to get batches of size 2. It prints each batch's data and labels.

PyTorch

import torch
from torch.utils.data import Dataset, DataLoader

class SimpleDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

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

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# Sample data
features = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
labels = torch.tensor([0, 1, 0, 1])

# Create dataset
dataset = SimpleDataset(features, labels)

# Use DataLoader to get batches
loader = DataLoader(dataset, batch_size=2, shuffle=True)

for batch_data, batch_labels in loader:
    print('Batch data:', batch_data)
    print('Batch labels:', batch_labels)
    print('---')

Important Notes

Always return data and label as tensors in __getitem__ for PyTorch models.

You can add data transformations inside __getitem__ if needed.

Using DataLoader with your custom dataset helps with batching and shuffling automatically.

Summary

A Custom Dataset class organizes your data for PyTorch models.

Implement __init__, __len__, and __getitem__ methods.

Use DataLoader with your dataset to handle batches and shuffling easily.

Practice

1. What is the main purpose of creating a Custom Dataset class in PyTorch?

easy

A. To visualize data samples

B. To organize and provide access to data samples for model training

C. To perform gradient descent optimization

D. To define the neural network architecture

Solution

1. Step 1: Understand Dataset role

   A Custom Dataset class is designed to organize data and provide easy access to samples during training.

2. Step 2: Differentiate from other components

   Neural network architecture and optimization are handled elsewhere, not by the Dataset class.

3. Final Answer:

   To organize and provide access to data samples for model training -> Option B

4. Quick Check:

   Dataset class = data access [OK]

Hint: Dataset class organizes data for training access [OK]

Common Mistakes:

• Confusing Dataset with model definition
• Thinking Dataset handles optimization
• Assuming Dataset visualizes data

2. Which method must be implemented in a PyTorch Custom Dataset class to return the number of samples?

easy

A. __getitem__

B. __init__

C. __len__

D. forward

Solution

1. Step 1: Recall Dataset methods

   __len__ returns the total number of samples in the dataset.

2. Step 2: Identify other methods' roles

   __getitem__ fetches a sample, __init__ initializes the dataset, forward is for models, not datasets.

3. Final Answer:

   __len__ -> Option C

4. Quick Check:

   Count samples = __len__ [OK]

Hint: Count samples with __len__ method [OK]

Common Mistakes:

• Using __getitem__ to get length
• Confusing __init__ with length method
• Using forward method in Dataset

3. Given this Custom Dataset code, what will print(len(dataset)) output?

import torch
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self):
        self.data = [10, 20, 30, 40]
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx]

dataset = MyDataset()

medium

A. 4

B. 3

C. 40

D. Error

Solution

1. Step 1: Check __len__ method

   __len__ returns length of self.data, which is a list of 4 elements.

2. Step 2: Evaluate print output

   Calling len(dataset) calls __len__, so output is 4.

3. Final Answer:

   4 -> Option A

4. Quick Check:

   len([10,20,30,40]) = 4 [OK]

Hint: len(dataset) calls __len__ returning data length [OK]

Common Mistakes:

• Confusing length with last element value
• Expecting __getitem__ output instead of length
• Assuming error without checking code

4. Identify the error in this Custom Dataset class:

from torch.utils.data import Dataset

class SampleDataset(Dataset):
    def __init__(self, data):
        self.data = data
    def __len__(self):
        return self.data
    def __getitem__(self, idx):
        return self.data[idx]

sample = SampleDataset([1,2,3])
print(len(sample))

medium

A. Missing __init__ method

B. __getitem__ returns wrong index

C. Dataset class not imported

D. __len__ returns data list instead of its length

Solution

1. Step 1: Analyze __len__ method

   __len__ returns self.data, which is a list, not an integer length.

2. Step 2: Understand len() requirement

   len() expects __len__ to return an integer. Returning a list causes a TypeError.

3. Final Answer:

   __len__ returns data list instead of its length -> Option D

4. Quick Check:

   __len__ must return int, not list [OK]

Hint: __len__ must return integer length [OK]

Common Mistakes:

• Returning data list instead of length in __len__
• Ignoring import statements
• Confusing __getitem__ with __len__

5. You want to create a Custom Dataset class that returns both input features and labels from two separate lists. Which implementation correctly achieves this?

hard

A. class MyDataset(Dataset):
def __init__(self, inputs, labels):
self.inputs = inputs
self.labels = labels
def __len__(self):
return len(self.inputs)
def __getitem__(self, idx):
return self.inputs[idx], self.labels[idx]

B. class MyDataset(Dataset):
def __init__(self, inputs, labels):
self.data = inputs + labels
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx]

C. class MyDataset(Dataset):
def __init__(self, inputs, labels):
self.inputs = inputs
self.labels = labels
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return self.inputs[idx]

D. class MyDataset(Dataset):
def __init__(self, inputs, labels):
self.inputs = inputs
self.labels = labels
def __len__(self):
return len(self.inputs)
def __getitem__(self, idx):
return self.labels[idx]

Solution

1. Step 1: Check __init__ and __len__ consistency

   The correct implementation stores inputs and labels separately and returns length of inputs in __len__, assuming matching lengths with labels.

2. Step 2: Verify __getitem__ returns both features and labels

   The correct __getitem__ returns a tuple (self.inputs[idx], self.labels[idx]) for supervised learning.

3. Final Answer:

   return self.inputs[idx], self.labels[idx] -> Option A

4. Quick Check:

   Return (input, label) tuple in __getitem__ [OK]

Hint: Return (input, label) tuple in __getitem__ [OK]

Common Mistakes:

• Combining inputs and labels into one list
• Returning only inputs or only labels
• Mismatched lengths in __len__