Memory leaks (when early break from iterator)

[shekhovt]

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

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.