Does creating a data loader inside another data loader in pytorch slow things down (during meta-learning)?
I was trying to create a data loader for meta-learning but got that my code is extremely slow and I can't figure out why. I am doing this because a set of data sets (so I need data loaders for them) is what is used in meta-learning.
I am wondering if it's because I have a collate function generating data loaders.
Here is the collate function that generates data loaders (and receives ALL the data sets):
class GetMetaBatch_NK_WayClassTask:

    def __init__(self, meta_batch_size, n_classes, k_shot, k_eval, shuffle=True, pin_memory=True, original=False, flatten=True):
        self.meta_batch_size = meta_batch_size
        self.n_classes = n_classes
        self.k_shot = k_shot
        self.k_eval = k_eval
        self.shuffle = shuffle
        self.pin_memory = pin_memory
        self.original = original
        self.flatten = flatten

    def __call__(self, all_datasets, verbose=False):
        NUM_WORKERS = 0 # no need to change
        get_data_loader = lambda data_set: iter(data.DataLoader(data_set, batch_size=self.k_shot+self.k_eval, shuffle=self.shuffle, num_workers=NUM_WORKERS, pin_memory=self.pin_memory))
        #assert( len(meta_set) == self.meta_batch_size*self.n_classes )
        # generate M N,K-way classification tasks
        batch_spt_x, batch_spt_y, batch_qry_x, batch_qry_y = [], [], [], []
        for m in range(self.meta_batch_size):
            n_indices = random.sample(range(0,len(all_datasets)), self.n_classes)
            # create N-way, K-shot task instance
            spt_x, spt_y, qry_x, qry_y = [], [], [], []
            for i,n in enumerate(n_indices):
                data_set_n = all_datasets[n]
                dataset_loader_n = get_data_loader(data_set_n) # get data set for class n
                data_x_n, data_y_n = next(dataset_loader_n) # get all data from current class 
                spt_x_n, qry_x_n = data_x_n[:self.k_shot], data_x_n[self.k_shot:] # [K, CHW], [K_eval, CHW]
                # get labels
                if self.original:
                    #spt_y_n = torch.tensor([n]).repeat(self.k_shot)
                    #qry_y_n = torch.tensor([n]).repeat(self.k_eval)
                    spt_y_n, qry_y_n = data_y_n[:self.k_shot], data_y_n[self.k_shot:]
                else:
                    spt_y_n = torch.tensor([i]).repeat(self.k_shot)
                    qry_y_n = torch.tensor([i]).repeat(self.k_eval)
                # form K-shot task for current label n
                spt_x.append(spt_x_n); spt_y.append(spt_y_n) # array length N with tensors size [K, CHW]
                qry_x.append(qry_x_n); qry_y.append(qry_y_n) # array length N with tensors size [K, CHW]
            # form N-way, K-shot task with tensor size [N,W, CHW]
            spt_x, spt_y, qry_x, qry_y = torch.stack(spt_x), torch.stack(spt_y), torch.stack(qry_x), torch.stack(qry_y)
            # form N-way, K-shot task with tensor size [N*W, CHW]
            if verbose:
                print(f'spt_x.size() = {spt_x.size()}')
                print(f'spt_y.size() = {spt_y.size()}')
                print(f'qry_x.size() = {qry_x.size()}')
                print(f'spt_y.size() = {qry_y.size()}')
                print()
            if self.flatten:
                CHW = qry_x.shape[-3:]
                spt_x, spt_y, qry_x, qry_y = spt_x.reshape(-1, *CHW), spt_y.reshape(-1), qry_x.reshape(-1, *CHW), qry_y.reshape(-1)
            ## append to N-way, K-shot task to meta-batch of tasks
            batch_spt_x.append(spt_x); batch_spt_y.append(spt_y)
            batch_qry_x.append(qry_x); batch_qry_y.append(qry_y)
        ## get a meta-set of M N-way, K-way classification tasks [M,K*N,C,H,W]
        batch_spt_x, batch_spt_y, batch_qry_x, batch_qry_y = torch.stack(batch_spt_x), torch.stack(batch_spt_y), torch.stack(batch_qry_x), torch.stack(batch_qry_y)
        return batch_spt_x, batch_spt_y, batch_qry_x, batch_qry_y

that is passed to another data loader here:
def get_meta_set_loader(meta_set, meta_batch_size, n_episodes, n_classes, k_shot, k_eval, pin_mem=True, n_workers=4):
    &quot;&quot;&quot;[summary]

    Args:
        meta_set ([type]): the meta-set
        meta_batch_size ([type]): [description]
        n_classes ([type]): [description]
        pin_mem (bool, optional): [Since returning cuda tensors in dataloaders is not recommended due to cuda subties with multithreading, instead set pin=True for fast transfering of the data to cuda]. Defaults to True.
        n_workers (int, optional): [description]. Defaults to 4.

    Returns:
        [type]: [description]
    &quot;&quot;&quot;
    if n_classes &gt; len(meta_set):
        raise ValueError(f'You really want a N larger than the # classes in the meta-set? n_classes, len(meta_set = {n_classes, len(meta_set)}')
    collator_nk_way = GetMetaBatch_NK_WayClassTask(meta_batch_size, n_classes, k_shot, k_eval)
    episodic_sampler = EpisodicSampler(total_classes=len(meta_set), n_episodes=n_episodes)
    episodic_metaloader = data.DataLoader(
        meta_set, 
        num_workers=n_workers, 
        pin_memory=pin_mem, # to make moving to cuda more efficient
        collate_fn=collator_nk_way, # does the collecting to return M N,K-shot task
        batch_sampler=episodic_sampler # for keeping track of the episode
        )
    return episodic_metaloader

(will generate a smaller example)

related:

https://discuss.pytorch.org/t/what-does-runtimeerror-cuda-driver-error-initialization-error-mean/87505/6


Conceptually pytorch dataloaders should have no problem being fast even if one is inside the other. One way to debug your issue is to use the line_profiler package to get a better idea of where the slowdown happens.
If you cannot resolve the issue after using the line_profiler, please update your questions with the output of the profiler to help us understand what might be wrong. Allow the profiler to run for some time to gather enough statistics about the execution of your dataloader. The @profile decorator works for both functions and class functions too so it should work for your dataloader functions.
