Memory issues and unloading data from memory

[mattwarkentin]

Hi @fepegar,

I'm not sure if this question/discussion makes much sense, but here goes. I've been running into some memory issues during training and I've been trying to think about what might be happening. My batch sizes are already pretty small (~2-6). I know torchio uses lazy loading to wait to load images until necessary, but as the epoch proceeds forward, isn't it true that more and more subjects will have their data loaded into memory? Is this memory every released?

Of course we only ever need the current batch in memory, and all previous subjects can be unloaded. I'm wondering if the Image class should have an unload() method to go along with the load() to relinquish memory to keep the memory footprint as small as possible. Similarly, the Subject class would have an unload() method to unload all images from memory.

Maybe I have an incorrect mental model of how things are working behind the scenes. I look forward to discussing.

[mattwarkentin]

I don't know if it is super hacky or bad form but something like:

class Image(dict):
    def unload(self):
        new_image = Image(...)
        self.__dict__.update(new_image.__dict__)

[mattwarkentin]

Or possibly something like:

class Image(dict):
    @classmethod
    def unload(cls):
        return cls(**self.__dict__)

[fepegar]

This is a good point! But the memory shouldn't grow thanks to the copy that is performed here:

torchio/torchio/data/dataset.py

Line 82 in db45b77

subject = copy.deepcopy(subject) # cheap since images not loaded yet

So

new unloaded image = copy(unloaded image)
loaded image = transform(new unloaded image)

and the original unloaded image doesn't add to the memory usage. Does this make sense?

[mattwarkentin]

Got it. Thanks for the explanation!

[mattwarkentin]

While we are discussing...it is nice that the loading and transforms are done on copies since it avoids memory explosions, but is there any interest in supporting cached subject datasets similar to those offered by monai? The data loading is the computational bottleneck for my current project, and I don't use any random transforms so in theory after the first epoch the data loading could be exponentially faster if transformed data were cached somehow. Thoughts?

[mattwarkentin]

I think monai does all the caching up front (actually it looks like PersistentDataset does real-time caching of non-random transforms), and then uses the cached data during training, it might be nice to have a CachedSubjectsDataset that caches subjects in real-time, and then the cached data is loaded upon next time trying to access that subject...I'm just thinking out loud

[fepegar]

Yes, I also like MONAI's caching features! What I don't like is that some of the subjects are not used in some of the epochs, I believe. I might be wrong.

What I've done some other times to "solve" this is running the preprocessing beforehand and saving the resulting images uncompressed. They are loaded pretty fast, especially on SSD. This is perhaps less elegant, but it has worked for me in the past. Unfortunately, it's not in my plans to work extensively on large new features within the next months, as I need to finish my thesis.

I'm sure there are other ways of optimizing TorchIO, for example avoiding some copies and conversions to ITK. Maybe we could have an in-place mode where no copies were ever made, that might make things much faster. Also, I don't know if you're using patches, but I think it wouldn't be hard to write a sampler that loads patches directly from memory using NiBabel.

[mattwarkentin]

Yes, I also like MONAI's caching features! What I don't like is that some of the subjects are not used in some of the epochs, I believe. I might be wrong.

Hmm, I think that is still covered by the sampler in the PyTorch DataLoader, so I think every subject would get used if you so desire.

What I've done some other times to "solve" this is running the preprocessing beforehand and saving the resulting images uncompressed.

I have done this too, but any time I want to change any little thing about the preprocess transforms, these images are now out of date. I like the idea of real-time processing that torchio offers, but it would be nice to have the benefit of caching in the situations where not much changes.

Unfortunately, it's not in my plans to work extensively on large new features within the next months, as I need to finish my thesis.

No problem! I also need to finish my thesis, so I can definitely understand. Perhaps when I am done my thesis work I will try to make an attempt at emulating something like monai.data.PersistentDataset but for torchio. I think it would be a fantastic addition and offer massive speed increases in most use cases.