Given a torch DataLoader, how can I create another DataLoader which is a subset

Viewed 92

Given an torch.utils.data.dataloader.DataLoader (let's say stored as variable data_loader_original), how can I create a new torch.utils.data.dataloader.DataLoader which contains a subset of data_loader_original?

I've seen this post how to adjust dataloader and make a new dataloader?, but I don't want to take a subset of the dataset directly (i.e. I don't want to use torch.utils.data.Subset). Instead I want something like

data_loader_subset = create_subset_data_loader(data_loader_original)

or

data_loader_subset = data_loader_ogirinal[:size_of_subset]

These are just examples to help illustrate what I'm looking for.

1 Answers

I'm not sure how can you split the subset, for the simple version, the snipcode below may help:

import torch
from torch.utils.data import DataLoader

bs = 50
shuffle = False
num_workers = 0
dataset = torch_dataset()
data_loader_original = DataLoader(dataset, batch_size=bs, shuffle=shuffle)

def create_subset_data_loader(loader, size_of_subset):
    count = 0
    for data in loader:
        yield data
        if count == size_of_subset:
            break
        count+=1

size_of_subset = 10

for epoch in range(epochs):
   for data in create_subset_data_loader(data_loader_original, size_of_subset):
      # processing

Related