Memory leaks (when early break from iterator)

[shekhovt]

It is eating up host memory indefinitely, much more than the unpacked dataset would take.

Duplicate of #345 , the explanation and a better workaround are proposed there.

MWE:

import torch
#------------------
import ffcv
from ffcv.loader import OrderOption
from ffcv.fields.decoders import IntDecoder
from ffcv.transforms import ToTensor, Squeeze, ToTorchImage
from ffcv.fields.rgb_image import CenterCropRGBImageDecoder
# -----------------

dev = torch.device('cpu')

train_set = '/local/temporary/data/imagenet/imagenet_pytorch/train-res=256-cls=10.beton'
input_size = 128

def train_pipelines():
    image_pipeline= [
            CenterCropRGBImageDecoder((input_size, input_size), ratio = 1),
            ToTensor(),
            ToTorchImage(),
        ]
    
    label_pipeline = [
            IntDecoder(),
            ToTensor(),
            Squeeze(),
        ]
    
    # Pipeline for each data field
    tr_pipelines = {
        'image': image_pipeline,
        'label': label_pipeline
    }
    return tr_pipelines 

num_workers = 4
batch_size = 128
train_loader = ffcv.loader.Loader(train_set, batch_size=batch_size, num_workers=num_workers, order=OrderOption.RANDOM, pipelines=train_pipelines(), drop_last=True)

epochs = 1000

print('single loader')
for e in range(epochs):
    x = 0
    for data, target in train_loader:
        x += data.to(dev).sum().cpu().detach().item()
        break
    print(e, x)

The problem is apparently with the break out of the loader loop. For some intermediate tasks we need just a few batches. There must be a way to cleanup?

[shekhovt]

I found a workaround to replace the loop part with this:

    it = iter(train_loader)
    for data, target in it:
        x += data.to(dev).sum().cpu().detach().item()
        break
    it.close()

I am not expert on python, but it appears that

    def __del__(self):
        self.close()

ffcv/ffcv/loader/epoch_iterator.py

Line 174 in 3a12966

def __del__(self):

is not a safe way to free resources.

[shekhovt]

If the image pipeline includes GPU processing, it is eating up GPU memory as well,

Some augmented tensors (already used or prepared ahead) are not disposed properly. This is in the context of MWE. The workaround stops that.

[richardrl]

Same issue, memory leak with early dataloader break. Testing your fix

[richardrl]

I'm not sure if this fixed things. I applied this and still get memory leaks. FFCV has serious memory leak issues

[shekhovt]

Frankly, I do not fully understand why the destructor may not be called. But the workaround I wrote above does work for me, I also made it a context manager to encapsulate the tidy-up work:

class SafeIter:
    def __init__(self, loader):
        self.loader = loader
        self.it = None

    def __enter__(self):
        self.it = iter(self.loader)
        return self.it
    
    def __exit__(self, exc_type, exc_value, exc_tb):
        if self.it is not None:
            if hasattr(self.it, 'close'):
                self.it.close()
            torch.cuda.empty_cache()

To be used as

with SafeIter(loader) as it:
            for i, (data, target) in enumerate(it):
                ...

Notice the empty_cache line. It is not needed, but depending how you monitor the memory usage you might observe that GPU memory is 'leaking' while in fact it is just being lazily freed. FFCV is also using a lot of file cache kind of memory. I believe this is by design and considered as free or reclaimable by the system. Maybe you could check that it is not locators caches and isolate the issue if you still observe leaks.

[richardrl]

@shekhovt Thank you for your response. Your solution is similar to #206

I was able to fix the memory leak using #206

It was also crucial for me to not use tqdm. I was using tqdm.tqdm(dataloader), and this still caused leaks. Using #206 + iter(train_dataloader) prevented the issue. Would be nice if there was a way to do things with TQDM though...