Memory leak

[shrubb]

Hi,

when the below example is run, the RAM usage grows forever:

import torch, torch.utils.data
import nonechucks

class DummyDataset(torch.utils.data.Dataset):
    def __len__(self):
        return 1_000_000

    def __getitem__(self, idx):
        return 666

dataset = nonechucks.SafeDataset(DummyDataset())

for _ in torch.utils.data.DataLoader(dataset):
    pass

Notes:

• Here the increase is quite slow; for a RAPID bug demonstration, replace 666 with torch.empty(10_000) (be careful to kill the process in time, before you're OOM!).
• No problems without SafeDataset.
• Without torch.utils.data.DataLoader, the leak is still there, although at a smaller scale, around 1 MB of RAM is lost per 30000-40000 __getitem__ calls.
• PyTorch 1.0.1, nonechucks 0.3.1.

[aronhoff]

SafeDataset.__getitem__ memoizes the dataset items, with no parameter to change this.

This makes the assumption that the dataset will always fit in the memory.
Remove the @memoize line, and the leak is gone.

[msamogh]

Thanks for raising the issue, @shrubb!

Yeah, @aronhoff is spot on. I hadn't really thought about this use-case. Parameterising the memoization seems like a good idea! Would you be want to raise a PR for that, @aronhoff?

[taehyunoh]

@aronhoff 's solution works like a charm! I had been struggled by this.
If possible, please update pip repository as well. I'd appreciate it!
Thanks

[msamogh]

Well, I've put it there for a reason, and it still ensures that you don't iterate through the entire dataset every time for most cases where the dataset isn't larger than the memory.

I think parameterising is the best solution for now.

[aronhoff]

I tried adding a bool attribute to the memoize class, that would make its __call__ skip the lookup. The idea is that you could use it as a property of the method, e.g. self.__getitem__.memoize = False. (c57e70c)

Unfortunately this does not work with multiprocessing. __getitem__ is not in the __dict__ of an object, so pickling does not save or restore it and its attributes. Putting it into __dict__ does not solve this.

I do not currently have more time to find a way. Solution could be in having a custom __setstate__ in SafeDataset, or perhaps using a custom metaclass for it. Or do it in the __getitem__ function directly. Either way, it seems to need more entanglement between memoize and the owner object.

Perhaps concerns should be separated completely. You could have a MemoizedDataset that passes through the values from another dataset which is potentially a SafeDataset, while memoizing them.

Keep in mind that the cache dicts are going to be different instances accross processes, so you would be replicating your dataset for each subprocess. This may not be what you want, but for small datasets, it could be not worth the effort of dealing with read-write shared memory.

And with regards to the dataset, ImageNet would optimistically take around 200 GB :)

[msamogh]

Thanks for the PR, @shrubb! I have been quite busy, and haven't gotten time to look at it. I'll take a look soon and get back to you.

[timonbimon]

Has this been resolved? :)
Nonechucks looks pretty useful, but it's very normal for us to have datasets that are much largern than the RAM, so with the memory leak this would be a no-go.

[aksg87]

@timonbimon @aronhoff

Any solution? I just faced a memory leak with this as well!

[Erotemic]

Another issue with memoizing the __getitem__ is that you assume asking for the same index will always result in the same value.

For complicated training cases, executing __getitem__ twice might return wildly different batch items. It is not safe to assume that an index corresponds 1:1 with a specific item. For instance, in one of my more involved datasets all even indexes correspond to positive examples and odd indexes are for negative examples, but where those examples are somewhat random because an index actually chooses from a pool of examples associated with that index. This allows me to balance my positive / negative training cases, but it very much breaks assumptions made in this library.

[enric1994]

Face the same issue. Any workarounds?