Efficient Subsetting Techniques for PyTorch Datasets in Machine Learning and Neural Networks

In machine learning, especially when training neural networks, we often deal with large datasets. However, for various reasons, you might want to work with a smaller subset of the data:

• Development and Testing: It can be faster and more efficient to experiment with a smaller subset during development or to test your model's performance on unseen data.
• Limited Resources: If you have limited computational resources like memory or processing power, a subset can be used for training.
• Data Exploration: You might want to focus on specific categories or examples within the data for exploration purposes.

Approaches for Creating Subsets in PyTorch

Here are common approaches to create subsets of a PyTorch dataset:

1. Slicing (Not Recommended):

2. Indexing with a List:

   You can create a list of indices corresponding to the desired subset and use it to index the dataset:

   import torch
   
   dataset = ...  # Your PyTorch dataset
   subset_indices = [5, 12, 37, ...]  # List of desired indices
   
   subset_data = dataset[subset_indices]
   subset_labels = dataset.labels[subset_indices]  # Assuming labels exist
   

   This approach is straightforward but can be memory-intensive for large datasets.

3. torch.utils.data.Subset Class:

   PyTorch provides the Subset class for creating efficient subsets without modifying the original dataset:

   from torch.utils.data import Subset
   
   full_dataset = ...  # Your PyTorch dataset
   subset_indices = ...  # List of desired indices
   
   subset = Subset(full_dataset, subset_indices)
   

   This is the preferred method as it creates a new dataset object that only fetches the relevant data points when used with a data loader.

4. Custom Subset Class (Advanced):

   For more complex logic in defining subsets, you can create a custom class that inherits from torch.utils.data.Dataset:

   import torch
   
   class MyCustomSubset(torch.utils.data.Dataset):
       def __init__(self, full_dataset, filter_func):
           self.full_dataset = full_dataset
           self.filter_func = filter_func
   
       def __len__(self):
           return sum(self.filter_func(i) for i in range(len(self.full_dataset)))
   
       def __getitem__(self, idx):
           filtered_idx = 0
           for i in range(len(self.full_dataset)):
               if self.filter_func(i):
                   if filtered_idx == idx:
                       return self.full_dataset[i]
                   filtered_idx += 1
           raise IndexError
   
   # Example usage:
   def filter_by_class(idx, target_class):
       return dataset.labels[idx] == target_class
   
   subset = MyCustomSubset(full_dataset, lambda idx: filter_by_class(idx, 3))
   

   This approach offers greater flexibility for defining criteria for the subset but requires more code.

Using Subsets with Data Loaders

Once you have a subset, you can use it with a data loader for efficient batching and data augmentation during training:

from torch.utils.data import DataLoader

subset_loader = DataLoader(subset, batch_size=32, shuffle=True)

import torch

dataset = torch.utils.data.TensorDataset(  # Sample dataset
    torch.rand(1000, 3),  # Features
    torch.randint(0, 5, (1000,))  # Labels (0-4)
)

subset_indices = [5, 12, 37, 892]  # Example indices

subset_data = dataset[subset_indices][0]  # Accessing features (tensor)
subset_labels = dataset[subset_indices][1]  # Accessing labels (tensor)

print(subset_data.shape)  # Output: torch.Size([4, 3])
print(subset_labels.shape)  # Output: torch.Size([4])

torch.utils.data.Subset Class (Recommended):

from torch.utils.data import Subset

full_dataset = torch.utils.data.TensorDataset(  # Sample dataset
    torch.rand(1000, 3),  # Features
    torch.randint(0, 5, (1000,))  # Labels (0-4)
)

subset_indices = [5, 12, 37, 892]  # Example indices

subset = Subset(full_dataset, subset_indices)

# Accessing data points within a data loader (assuming data loader is created)
for data, label in subset:
    # Process data and label tensors
    pass

import torch

class MyCustomSubset(torch.utils.data.Dataset):
    def __init__(self, full_dataset, filter_func):
        self.full_dataset = full_dataset
        self.filter_func = filter_func

    def __len__(self):
        return sum(self.filter_func(i) for i in range(len(self.full_dataset)))

    def __getitem__(self, idx):
        filtered_idx = 0
        for i in range(len(self.full_dataset)):
            if self.filter_func(i):
                if filtered_idx == idx:
                    return self.full_dataset[i]
                filtered_idx += 1
        raise IndexError

# Example usage:
def filter_by_class(idx, target_class):
    return full_dataset.labels[idx] == target_class

full_dataset = torch.utils.data.TensorDataset(  # Sample dataset
    torch.rand(1000, 3),  # Features
    torch.randint(0, 5, (1000,))  # Labels (0-4)
)

subset = MyCustomSubset(full_dataset, lambda idx: filter_by_class(idx, 3))

# Accessing data points within a data loader (assuming data loader is created)
for data, label in subset:
    # Process data and label tensors
    pass

This approach might be suitable for smaller datasets where memory efficiency isn't a major concern. You can iterate through the original dataset and create a new list or another dataset object containing elements that meet your criteria:

import torch

dataset = ...  # Your PyTorch dataset

def filter_by_condition(data, label):
    # Define your filtering condition here (e.g., label == 2)
    return condition

subset_data = []
subset_labels = []
for datapoint, label in dataset:
    if filter_by_condition(datapoint, label):
        subset_data.append(datapoint)
        subset_labels.append(label)

# You can then convert these lists into a new dataset if needed

Third-Party Libraries (Conditional Samplers):

Libraries like torchsampler () offer functionalities for creating custom samplers that control how data points are selected during training. This allows for more intricate sampling logic within a data loader:

import torch
from torchsampler import ImbalancedDatasetSampler

# Assuming your dataset has imbalanced classes
class_counts = ...  # Get class counts from your dataset

sampler = ImbalancedDatasetSampler(dataset, class_counts)
subset_loader = DataLoader(dataset, sampler=sampler, batch_size=32)

Data Augmentation Libraries (Augmentation-Specific Subsets):

Libraries like albumentations () and imgaug () often provide functionalities to create subsets specifically for data augmentation purposes. These subsets might include transformations applied to the data points:

import torch
from albumentations import Compose, RandomHorizontalFlip

# Assuming your dataset contains images
aug_transforms = Compose([RandomHorizontalFlip()])

subset = MyDataset(data_paths, labels, transforms=aug_transforms)
# Use this subset for data augmentation during training

Choosing the Right Method:

The best method for creating subsets depends on your specific needs and dataset characteristics. Here's a general guideline:

• For small datasets and simple filtering, a loop-based approach might suffice.
• For memory efficiency and standard subsetting, use torch.utils.data.Subset.
• For complex filtering logic or imbalanced datasets, explore conditional samplers.
• For data augmentation workflows, leverage functionalities from data augmentation libraries.