PyTorch - Loading Data

PyTorch - Loading Data

Efficient data loading is crucial for training deep learning models. PyTorch provides powerful and flexible tools for this purpose through its torch.utils.data module, primarily the Dataset and DataLoader classes. The torchvision package also offers convenient pre-built datasets and image transformation utilities for computer vision tasks.

The Dataset Class

The torch.utils.data.Dataset is an abstract class representing a dataset. Your custom datasets should inherit Dataset and override two key methods:

• __len__(self): so that len(dataset) returns the size of the dataset.
• __getitem__(self, idx): to support indexing such that dataset[i] can be used to get the i-th sample.

Here’s a basic skeleton for a custom Dataset:

import torch
from torch.utils.data import Dataset

class MyCustomDataset(Dataset):
    def __init__(self, data_source, targets, transform=None):
        self.data = data_source # e.g., list of file paths, or pre-loaded data
        self.targets = targets
        self.transform = transform

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

    def __getitem__(self, idx):
        sample_data = self.data[idx] # Load or access your data point
        sample_target = self.targets[idx]

        # Apply transformations if any
        if self.transform:
            sample_data = self.transform(sample_data)

        return sample_data, sample_target

torchvision.datasets provides several pre-loaded datasets like CIFAR10, MNIST, etc., which are subclasses of torch.utils.data.Dataset.

For example, to use the CIFAR10 dataset:

import torchvision
import torchvision.transforms as transforms

# Define transformations: convert PIL Image to Tensor, then normalize
transform_pipeline = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)) # CIFAR10 mean/std for 3 channels
])

# Load CIFAR10 training set
trainset = torchvision.datasets.CIFAR10(root='./data',
                                        train=True,
                                        download=True,
                                        transform=transform_pipeline)

# Load CIFAR10 test set
testset = torchvision.datasets.CIFAR10(root='./data',
                                       train=False,
                                       download=True,
                                       transform=transform_pipeline)

transforms.ToTensor() converts a PIL Image or NumPy ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]. transforms.Normalize() normalizes the tensor image with a given mean and standard deviation for each channel.

The DataLoader Class

The torch.utils.data.DataLoader is an iterator that provides many essential features for loading data:

• Batching the data.
• Shuffling the data at every epoch.
• Loading data in parallel using multiprocessing workers.

import torch

# Assuming trainset is an instance of a Dataset (e.g., the CIFAR10 trainset from above)
trainloader = torch.utils.data.DataLoader(trainset,
                                          batch_size=64, # Number of samples per batch
                                          shuffle=True,    # Shuffle data at every epoch
                                          num_workers=2)   # Number of subprocesses for data loading

# Similarly for the test set, usually without shuffling
testloader = torch.utils.data.DataLoader(testset,
                                         batch_size=100,
                                         shuffle=False,
                                         num_workers=2)

batch_size determines how many samples are processed before the model’s weights are updated. shuffle=True is important for training to prevent the model from learning the order of samples. num_workers > 0 enables multi-process data loading, which can significantly speed up training by preparing data in parallel on the CPU while the GPU is busy with computations. A common practice is to start with num_workers=0 (for easier debugging, as data loading occurs in the main process) and then increase it, e.g., to os.cpu_count(), to find an optimal value based on system performance and specific workload.

Example: Custom Dataset for CSV Files

Often, your data might reside in a CSV file. You can use libraries like Pandas to read the CSV and then wrap it in a custom Dataset.

Consider a CSV file named my_data.csv with features and labels:

Example my_data.csv content: feature1,feature2,label 1.5,2.3,0 4.1,3.9,1 0.8,1.1,0

import pandas as pd
import torch
from torch.utils.data import Dataset

class CSVTabularDataset(Dataset):
    def __init__(self, csv_path, transform=None):
        self.dataframe = pd.read_csv(csv_path)
        # Assuming the last column is the target/label, and others are features
        self.features = self.dataframe.iloc[:, :-1].values.astype('float32')
        self.labels = self.dataframe.iloc[:, -1].values.astype('int64') # Labels often int64 for CrossEntropyLoss
        self.transform = transform

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

    def __getitem__(self, idx):
        feature_sample = torch.tensor(self.features[idx], dtype=torch.float32)
        label_sample = torch.tensor(self.labels[idx], dtype=torch.long)

        if self.transform:
            # Example: transform could be a normalization function for features
            feature_sample = self.transform(feature_sample)

        return feature_sample, label_sample

# Usage (assuming 'my_data.csv' exists or is created):
# # To create a dummy my_data.csv for this example:
# # with open('my_data.csv', 'w') as f:
# #     f.write('feature1,feature2,label\n1.5,2.3,0\n4.1,3.9,1\n0.8,1.1,0')
#
# tabular_dataset = CSVTabularDataset(csv_path='my_data.csv')
# tabular_loader = torch.utils.data.DataLoader(tabular_dataset, batch_size=2, shuffle=True)
#
# # Iterate through the DataLoader
# # for features, labels in tabular_loader:
# #     print(f'Batch Features Shape: {features.shape}, Batch Labels: {labels}')
# #     # Example: Batch Features Shape: torch.Size([2, 2]), Batch Labels: tensor([0, 1])
# #     break

Handling Complex CSV Data (e.g., Image Landmarks)

If your CSV contains more complex data, like file paths to images and associated landmark coordinates, your __getitem__ method will involve more steps, such as loading images from paths and parsing coordinate strings. For instance, updating the original tutorial’s landmark processing logic:

import pandas as pd
import numpy as np # Or use torch for reshaping/tensor conversion
import torch
# from PIL import Image # For loading images

# This snippet shows logic that would go inside a custom Dataset's __getitem__ method.
# Assume 'self.dataframe' is a pandas DataFrame loaded from a CSV like 'faces/face_landmarks.csv'.
# Each row might contain: image_filename, landmark1_x, landmark1_y, landmark2_x, ...
# And 'idx' is the index passed to __getitem__.

# Inside a custom Dataset's __getitem__(self, idx):
# current_row = self.dataframe.iloc[idx]
# img_path = current_row.iloc[0] # First element, e.g., 'path/to/image_001.jpg'
# landmark_values_str = current_row.iloc[1:].values # Rest are landmark coordinate strings

# Convert landmark coordinate strings (read from CSV) to a numeric array and reshape
# landmarks_numeric = landmark_values_str.astype('float32').reshape(-1, 2) # Assuming 2D landmarks (x, y)
# landmarks_tensor = torch.tensor(landmarks_numeric, dtype=torch.float32)

# You would then typically load the image using img_path:
# image = Image.open(img_path).convert('RGB') # Convert to RGB for consistency

# Apply necessary transforms to both image and landmarks (if any).
# if self.transform:
#    image, landmarks_tensor = self.transform(image, landmarks_tensor) # Transform might need to handle both

# return image, landmarks_tensor

# Note: The original snippet:
# landmarks_frame = pd.read_csv('faces/face_landmarks.csv')
# n = 65
# img_name = landmarks_frame.iloc[n, 0]
# landmarks = landmarks_frame.iloc[n, 1:].values # .values replaces .as_matrix()
# landmarks = landmarks.astype('float').reshape(-1, 2)
# This logic is for extracting data for a single, specific row 'n' from the entire dataframe.
# In a Dataset, you'd adapt this for the current 'idx'.

Common Transformations (torchvision.transforms)

torchvision.transforms provides a suite of common image transformations. These can be chained together using transforms.Compose. Some frequently used transforms include:

• transforms.ToTensor(): Converts a PIL Image or NumPy ndarray to a torch.Tensor and scales values to [0,1].

• transforms.ToPILImage(): Converts a torch.Tensor or ndarray to a PIL Image.

• transforms.Normalize(mean, std): Normalizes a tensor image with mean and standard deviation per channel.

• transforms.Resize(size): Resizes the input image to the given size.

• transforms.CenterCrop(size): Crops the given image at the center.

• transforms.RandomResizedCrop(size): Crops the image to a random size and aspect ratio, then resizes to the target size.

• transforms.RandomHorizontalFlip(p=0.5): Horizontally flips the given image randomly with a probability p (default 0.5).

• Data augmentation transforms like transforms.ColorJitter, transforms.RandomRotation, transforms.RandomAffine, etc., help improve model robustness by exposing it to a wider variety of training data.

Proper data augmentation and preprocessing are key to training robust models. For more details, refer to the torchvision.transforms documentation.

Practical Considerations and Best Practices

• Debugging DataLoaders: Start with num_workers=0. This loads data in the main process, making debugging easier (e.g., stack traces are clearer). Once confirmed, increase num_workers for performance.

• Batch Size: Choose a batch_size that fits your GPU memory and provides stable gradients. Powers of 2 (e.g., 32, 64, 128, 256) are common due to hardware optimizations, but the optimal value is empirical.

• Reproducibility: For reproducible results, especially when shuffle=True or using random transforms, set random seeds for PyTorch (torch.manual_seed()), NumPy (np.random.seed()), and Python’s random module. For DataLoader with num_workers > 0, you might also need a custom worker_init_fn to ensure workers also have their seeds set correctly if they use randomness independently.

• Large Datasets: For datasets too large to fit in memory, ensure your Dataset loads data on-the-fly in __getitem__ (e.g., reads an image from disk). For truly massive datasets or data streaming scenarios, consider torch.utils.data.IterableDataset, which is suitable for cases where random access is not feasible or data is generated dynamically. Libraries like WebDataset can also be helpful for large-scale distributed training.

• Collate Function: DataLoader uses a collate_fn to combine individual samples (returned by Dataset.__getitem__) into a batch. The default collate_fn works well for many standard data types (tensors, numbers, strings). However, for complex data structures like sequences of varying lengths or custom objects, you may need to provide a custom collate_fn to specify how samples should be batched together (e.g., padding sequences to the same length).

Real-World Application: Image Classification

In an image classification task, you might have a folder structure where each subfolder represents a class (e.g., /data/cats/, /data/dogs/). Your custom Dataset would:

1. Scan these folders to find all image paths and assign corresponding labels (e.g., 0 for cat, 1 for dog).
2. In __getitem__, given an index, load the image file from its path (e.g., using PIL.Image.open()).
3. Apply torchvision.transforms: resize to a standard input size, perhaps perform random augmentations (like flips or rotations), convert the image to a torch.Tensor via transforms.ToTensor(), and then normalize it using transforms.Normalize() with pre-calculated mean and standard deviation for the dataset.

The DataLoader then takes this dataset, shuffles the data (for training), and creates batches of (image_tensor_batch, label_tensor_batch). These batches are fed sequentially into your neural network for training. This efficient data pipeline is crucial for preventing the GPU from idling and thus maximizing training throughput.