Memory leak while using enumerate() to deal with DataLoader provided by PyG

[zzzbbs]

I met a question while I training on the dataset while I do enumerate() to deal with DataLoader. The memory leaked slowly while doing enumerate, and finally stop training due to memory leak. I changed num_workers to zero, but it still happened. Is this a normal phenomenon while using enumerate() to deal with DataLoader? I will be very grateful if you can help me solve this problem.

[zzzbbs]

torch version 1.9.0, PyG version 2.2.0

[wsad1]

I doubt enumerate is causing this "memory leak". Could you please share the code and error you are seeing here.

[rusty1s]

This is not a dataloader issue but rather a dataset issue, right? It seems like you are running OOM to load in all your data into an InMemoryDataset. You could also use the Dataset interface which does not have the restriction to hold all data in RAM.

[zzzbbs]

This is not a dataloader issue but rather a dataset issue, right? It seems like you are running OOM to load in all your data into an InMemoryDataset. You could also use the Dataset interface which does not have the restriction to hold all data in RAM.

Thanks for your answer! Following your advice, I change my InMemoryDateset to Dataset. And loaded by DataLoader(train_set.get(1), batch_size=cfg.HYPER.BATCH_SIZE, shuffle=True, num_workers=8, pin_memory=False,drop_last=True) , the RAM still growing up slowly with batch executing untill memory leaked to swap space. The following is the code after changes. The preprocessed datasets are some 3D pose estimation dataset.

class NewDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(NewDataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        data_path = os.path.join(self.root, 'h5')
        self._raw_file_names = [os.path.join(data_path, file) for file in os.listdir(data_path)]
        return self._raw_file_names

    @property
    def processed_file_names(self):
        return ['data.pt']

    def process(self):
        data_list = []
        for file in self.raw_file_names:
            data = parse_data_numpy(file) #处理数据
            data_list.extend(data)
            
            if self.pre_filter is not None and not self.pre_filter(data):
                continue

            if self.pre_transform is not None:
                data = self.pre_transform(data)

        torch.save(data_list, osp.join(self.processed_dir, f'data.pt'))
    def len(self):
        return len(self.processed_file_names)            
    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, f'data.pt'))
        return data

[zzzbbs]

And where do you run into an OOM. Like which functions runs outs of memory?

After preprocessing data, I simply run the following code to enumerate dataset and print batch index, the RAM grows up slowly with batch executing untill memory leaked to swap space.
for batch_idx, data_list in enumerate(dataloader): print(batch_idx)

[wsad1]

In your get function you still seem to be loading the full dataset

def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, f'data.pt'))
        return data

In the above function use idx to load just a subset of the data. This tutorial on working with "large" datasets should help.

[zzzbbs]

In your get function you still seem to be loading the full dataset

def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, f'data.pt'))
        return data

In the above function use idx to load just a subset of the data. This tutorial on working with "large" datasets should help.

Thanks for your help! I will try to solve the question by dividing dataset into small subset.