Running out of memory processing dataset
I’m usingtrain_text_to_image.pyto train on a fairly large dataset of 920k images. During the preprocessing of the dataset everything seems okay until I get to “Extracting data files”, at which point it hangs:`Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:02<00:00, 343404.14it/s]Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:02<00:00, 339180.98it/s]Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:04<00:00, 205028.97it/s]Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:02<00:00, 348870.18it/s]Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:02<00:00, 364106.49it/s]Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:02<00:00, 374222.79it/s]Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:04<00:00, 219925.98it/s]Resolving data files: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:07<00:00, 127850.77it/s]Downloading and preparing dataset imagefolder/default to /drive/storage/data/REDACTED_TOKEN/imagefolder/default-6ce0ff74ab1e4790/0.0.0/37fbb85cc714a338bea574ac6c7d0b5be5aff46c1862c1989b20e0771199e93f…Downloading data files: 100%|████████████████████████████████████████████████████████████████████████████████████████| 920251/920251 [00:10<00:00, 84143.52it/s]Downloading data files: 0it [00:00, ?it/s]Extracting data files: 0it [00:00, ?it/s]If I monitor the process intopI can see that, when it gets to this point, a single python process runs at 100% with continuously increasing resident memory until it hits the system’s maximum. It will stay there for a while but eventually quit (possibly running out of swap, but I’m not sure). Earlier stages of processing run on 8 processes.I’m running using Accelerate on an 8x A100 80GB Google compute node with 1.3TB main memory. My initial batch size was 8 but I’ve also tried with 1 (I don’t think this has anything to do with batches, but I thought I’d give it a try).Any help greatly appreciated as I’m kind of stuck here!
2023-03-27T16:46:10Z
I see that there’s a “streaming” option toload_dataset. Would it makes sense to make thatTrue? And will that prevent shuffling the data?
2023-03-27T17:03:57Z
see:Streami had a similar issue where preprocessing of the dataset would just fill up the memory and oom’d. i solved it by using IterableDataset but i got the feeling it wasn’t desirable. i feel like preprocessing (specifically .map() fn) is for small, not memory-intensive operations like tokenization, and not for loading up large datasets into memory like images. i remember doing something hacky like using transform/augmentation to load an image on the fly while only storing the path but all of this is my limited experience, because .map() supposedly shouldn’t lead to oom because it doesn’t load dataset all at once. i just got frustrated at some point and decided to not figure out the right way.you can shuffle an iterable/streaming dataset, see above link, also look into trainer callbacks where you can invoke a reshuffle after each epoch. i haven’t tested below code but something like this should work:class ShuffleCallback(TrainerCallback):
    def on_epoch_begin(self, args, state, control, train_dataloader, **kwargs):
        if isinstance(train_dataloader.dataset, IterableDataset):
            train_dataloader.dataset.set_epoch(train_dataloader.dataset._epoch + 1)…andtrainer_object = Trainer(
...
callbacks=[ShuffleCallback()],
)
2023-03-27T17:30:13Z
Okay, I haven’t triedIterableDataset, as that will take a bit more work, but just settingstreaming=Truein theload_datacall doesn’t solve it. Still quits, just more quickly (which is kinda good, I guess).RegardingIterableDataset, is there a relatively quick way of trying that within the context oftrain_text_to_image.py? It usesImageFolderby default, and my data is prepared for that approach.UPDATE: Hmm… Actually, I see from the comments that setting thestreamingflag returns anIterableDataset… bummer. I’m stumped.
2023-03-27T17:43:21Z
are you still getting oom? with iterabledataset you shouldn’t get oom with a small enough batch size (try 2, although your setup should be able to handle much more).
2023-03-27T17:46:01Z
Well, the problem is it never returns fromload_dataset, so I don’t actuallygettheIterableDataset. It seems to be some processing happening insideload_datasetthat tops out the memory. I’m sure it’s clear, but this is not a CUDA out of memory error, this is system memory/RAM. I don’t get far enough for a CUDA memory error.PS - I tried a smaller batch size earlier, but I can give it another shot. But I do think the problem is prior to any kind of batching.
2023-03-27T18:08:04Z
jbmaxwell:train_text_to_image.pycan you try this?load_dataset(..., writer_batch_size=1)
2023-03-27T18:56:18Z
It hasn’t actually crashed yet, but it will soon:top - 18:54:25 up 13:16,  3 users,  load average: 1.01, 1.12, 2.51
Tasks: 758 total,   2 running, 756 sleeping,   0 stopped,   0 zombie
%Cpu(s):  0.6 us,  1.2 sy,  0.0 ni, 98.2 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
MiB Mem : 1370757.+total,  96039.5 free, 1271344.+used,   3373.6 buff/cache
MiB Swap:      0.0 total,      0.0 free,      0.0 used.  92884.0 avail Mem 

    PID USER      PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND                                                                                   
  69994 james     20   0 1348.8g   1.2t 246756 R 100.0  88.9  27:44.66 pythonIn case it’s useful:torch: 1.13.1+cu117diffusers: 0.15.0.dev0
2023-03-27T19:37:53Z
Maybe I’ll grab a localdiffuserscopy and make an editable install and see if I can figure out exactly where inload_datasetthings are going awry.
2023-03-29T09:45:16Z
Hi@jbmaxwell! Is it a public dataset so we can try to replicate? If it is not, can you try to replicate with another dataset and see if you still have the same issue? Also copying@lhoestqin case I’m missing something obvious here.
2023-03-29T14:31:22Z
It’s a private dataset, but it turned out there were missing columns names. Strange that it wasn’t easier to catch (and that the result was this memory leak problem), but once I found it everything went as expected. I had been trying to trace the problem fromload_databefore I realized what was wrong, but I never found the source of the memory issue.