CPU Memory exhausted when loading data with random_grid_cell_assignment()

[trettelbach]

Hi,
I'm currently loading data with the random_grid_cell_assignment(self.dataset_full, [0.7, 0.15, 0.15], grid_size=6). I had num_workers=4.
The dataset I was loading consisted of 3 files, each 15-16 GB in size; ca. 47 GB total.

When I ran the code, and monitored my resources with top, I saw that each individual process had loaded 47+ GB into its memory; cumulatively crashing my workstation as it ran out of CPU memory.

Is there a way that not all datasets are loaded entirely with all processes?

Thanks!

[adamjstewart]

A couple comments:

1. I don't think this has anything to do with random_grid_cell_assignment but it would be good to see if you can reproduce this bug without using it to narrow down the issue.
2. What is the output of gdalinfo on one of your files? By default, TorchGeo uses rasterio windowed reading so that we only load a single block of data at a time instead of loading the entire file into memory every time. I'm guessing your files are not sharded, or the block size is too large. I would recommend converting them to COGs like so:

   > gdal_translate -of COG input.tif output.tif

[trettelbach]

Thanks for the quick reply! So I did run it without the random_grid_cell_assignment but instead declared the train, val, test sets like so:

self.dataset_train = Image(f'{self.data_dir}/train/') & Mask(f'{self.data_dir}/')
self.dataset_valid = Image(f'{self.data_dir}/valid/') & Mask(f'{self.data_dir}/')
self.dataset_test = Image(f'{self.data_dir}/test/') & Mask(f'{self.data_dir}/')

where ca. 200 GB of files where in {self.data_dir}/train and this did not overload the CPU memory - which originally lead me to assume that the random_grid_cell_assignment was the culprit.

The gdalinfo output of one of the files is:

Driver: GTiff/GeoTIFF
Files: filename.tif
Size is 80888, 23500
Coordinate System is:
PROJCRS["WGS 84 / UTM zone 4N",
    BASEGEOGCRS["WGS 84",
        DATUM["World Geodetic System 1984",
            ELLIPSOID["WGS 84",6378137,298.257223563,
                LENGTHUNIT["metre",1]]],
        PRIMEM["Greenwich",0,
            ANGLEUNIT["degree",0.0174532925199433]],
        ID["EPSG",4326]],
    CONVERSION["UTM zone 4N",
        METHOD["Transverse Mercator",
            ID["EPSG",9807]],
        PARAMETER["Latitude of natural origin",0,
            ANGLEUNIT["degree",0.0174532925199433],
            ID["EPSG",8801]],
        PARAMETER["Longitude of natural origin",-159,
            ANGLEUNIT["degree",0.0174532925199433],
            ID["EPSG",8802]],
        PARAMETER["Scale factor at natural origin",0.9996,
            SCALEUNIT["unity",1],
            ID["EPSG",8805]],
        PARAMETER["False easting",500000,
            LENGTHUNIT["metre",1],
            ID["EPSG",8806]],
        PARAMETER["False northing",0,
            LENGTHUNIT["metre",1],
            ID["EPSG",8807]]],
    CS[Cartesian,2],
        AXIS["(E)",east,
            ORDER[1],
            LENGTHUNIT["metre",1]],
        AXIS["(N)",north,
            ORDER[2],
            LENGTHUNIT["metre",1]],
    USAGE[
        SCOPE["Engineering survey, topographic mapping."],
        AREA["Between 162°W and 156°W, northern hemisphere between equator and 84°N, onshore and offshore. United States (USA) - Alaska (AK)."],
        BBOX[0,-162,84,-156]],
    ID["EPSG",32604]]
Data axis to CRS axis mapping: 1,2
Origin = (580403.734170000068843,7758041.552710000425577)
Pixel Size = (0.100000000000000,-0.100000000000000)
Metadata:
  AREA_OR_POINT=Area
Image Structure Metadata:
  INTERLEAVE=PIXEL
Corner Coordinates:
Upper Left  (  580403.734, 7758041.553) (156d54' 5.59"W, 69d55' 2.53"N)
Lower Left  (  580403.734, 7755691.553) (156d54'13.17"W, 69d53'46.72"N)
Upper Right (  588492.534, 7758041.553) (156d41'26.33"W, 69d54'53.10"N)
Lower Right (  588492.534, 7755691.553) (156d41'34.68"W, 69d53'37.30"N)
Center      (  584448.134, 7756866.553) (156d47'49.92"W, 69d54'20.03"N)
Band 1 Block=80888x1 Type=UInt16, ColorInterp=Gray
  NoData Value=0
Band 2 Block=80888x1 Type=UInt16, ColorInterp=Undefined
  NoData Value=0
Band 3 Block=80888x1 Type=UInt16, ColorInterp=Undefined
  NoData Value=0
Band 4 Block=80888x1 Type=UInt16, ColorInterp=Undefined
  NoData Value=0

I will try how it works as a COG. Thanks for the tip!

[adamjstewart]

Ah, I'm guessing that each directory only contains a single image and that's why it doesn't run out of memory? I'll bet if you use a single dataset containing all images for train, val, and test it will run out of memory without random_grid_cell_assignment.

[trettelbach]

No actually, train/ contains 15 files adding up to ca. 200 GB.
When running random_grid_cell_assignment I pointed it to that same directory (but used regex to only read 3 of those files adding up to 45 GB).

But another user of the workstation has also encountered weird behavior of the system. I'll make sure to exclude the hardware as a source of the error and try running it again.

[trettelbach]

Alright, so I have converted all my datasets to COGs and reran the training.
The good thing is: Using COGs did reduce the overall RAM need! The bad thing: I was still experiencing sort of a memory leak (not in the strict sense probably) throughout the epochs - the run just crashed later than when I wasn't using COGs.
This is what my CPU memory usage graph looked like (training started at 11.30 - OOM in epoch 31):

In the meantime, I have found the following issues:
pytorch/pytorch#13246
ultralytics/ultralytics#1467 (comment)
It looks like there is a CPU memory problem with the DataLoader when num_workers > 0.

Running my training with num_workers = 0 and on multiple GPUs, I did succeed to train for my desired 50 epochs; however, training was very slow at >60 mins/epoch.

What I have also found, is that using strategy='ddp' in the Trainer for multi-GPU training would vastly increase my CPU memory needs.
I have therefore now returned to using only a single GPU and have cautiously set num_workers=4 (resulting in ca. 8 mins/epoch). My CPU memory needs are a lot lower, yet still slowly increasing. I have reduced my overall training to 30 epochs. This is what my CPU memory usage looks like now (I manually killed python at 17:52):

I manage to get in all my training before I run out of memory, however, there is something happening afterwards, that I cannot explain (I am not predicting, just running trainer.fit()).
For anybody who might encounter similar behaviour: My current workaround is to sys.exit(1) after the last epoch, which avoids whatever is happening with the CPU memory afterwards.

But overall, it seems like the DataLoader and num_workers problem is a known pytorch issue that doesn't have an official fix at this point. Just thought I'd share what I've found out and what my workarounds look like, in case someone else ends up here with similar problems.

[adamjstewart]

Can you try modifying this line to make it smaller and see what happens? This caches the 128 most recent filehandles, maybe that's too many.