CPU RAM Memory usage

[patriksabol]

Hello guys,

I am training model on Chesapeake13 labels and ChesapeakeCVPR rasters, from which I pull only naip-new.tif rasters off.
I use RandomBatchGeoSampler using pathsize=256. When cache=True for DataLoader, during training, the CPU RAM memory is constantly increasing. I have found length=50000 is just to fit into 200GB RAM memory. Using higher length leads to CPU RAM memory overflow. Is this expected behaviour?

[adamjstewart]

Is this on 0.4.1 or on the main branch? If the latter, it could be related to #1331.

[adamjstewart]

In that case I'm not sure what the issue is. @yichiac reported similar issues.

[patriksabol]

After some testing, I have found out that this issue is related to cache option in RasterDataset. I was running test of the model on Chesapeake labels and NAIP CVPR Imagery. When cache=True for both datasets (Chesapeake13, NAIP CVPR Imagery), the process is draining CPU RAM Memory very quickly.

When cache=False, CPU RAM memory is stable

but it is significantly slower, about 10x slower.
Is it done by design? Is there any way, how to limit cache memory?

[adamjstewart]

If I recall, there are actually 2 kinds of caches going on that may be causing your issue.

In TorchGeo, we cache the file handles of the datasets we load: https://github.com/microsoft/torchgeo/blob/v0.4.1/torchgeo/datasets/geo.py#L488

When you use cache=False, it disables that. We only save up to 128 file handles. I find it surprising that this eats up so much memory.

I see a couple of possibilities. One is that the image size or resolution you're using results in very large images. I think this is unlikely. The other possibility is that you're being bitten by GDAL's internal cache. If the files you are loading have a very large block size (or no blocks at all) then that could explain why saving only 128 file handles is enough to run out of memory.

Can you try setting GDAL_CACHEMAX to a small value and see if that helps? See https://gdal.org/user/configoptions.html#performance-and-caching for details, there are a few other settings you could play around with in there.

@calebrob6 have you run into any issues like this using our Chesapeake datasets?

[patriksabol]

The 'Chesapeake13' label tif has a block size of 128x128, while the NAIP Imagery from the CVPR zip file has a block size of 512x512. Therefore, there should not be any issues.

For the test split of the CVPR NAIP Imagery, I discovered that setting GDAL_CACHEMAX=1% (default is 5%) allowed the test phase to run successfully without compromising performance. This was achieved while 16GB of CPU RAM was available.

But it still depends on the size of used dataset and CPU RAM Memory.
For training, with default settings and 200GB CPU RAM memory, I can use only 50 000 training samples, which consume the whole available RAM memory.

It seems like @functools.lru_cache(maxsize=128) is not functioning as intended.

[adamjstewart]

This does seem like an issue. Can you open a formal issue with a minimal reproducible example? I'm trying to complete and defend my thesis this month so I can't guarantee I'll get to it until August but this is something I definitely want to fix. We saw issues similar to this in our testing so I want to better understand what's going on and squash this bug before the next release if we can.