[Python] Memory keeps increasing in pytorch training loop, even with empty_cache()

I have a pytorch training script, and I'm getting an out-of-memory error after a few epochs even tho I'm calling torch.cuda.empty_cache(). The GPU memory just keeps going up and I can't figure out why.

Here's basically what I'm doing:

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

class CustomDataset(Dataset):
def __init__(self, data_paths):
self.data_paths = data_paths

def __len__(self):
return len(self.data_paths)

def __getitem__(self, idx):
image = np.load(self.data_paths[idx]['image']).astype(np.float32)
label = np.load(self.data_paths[idx]['label']).astype(np.int64)

image = torch.tensor(image).cuda()
label = torch.tensor(label).cuda()

return image, label

data_paths = [{'image': f'img_{i}.npy', 'label': f'label_{i}.npy'} for i in range(10000)]
dataset = CustomDataset(data_paths)
dataloader = DataLoader(dataset, batch_size=32, num_workers=4, pin_memory=True)

for epoch in range(10):
for batch in dataloader:
images, labels = batch

output = images.mean()
loss = output.sum()
loss.backward()

del images, labels, loss, output
torch.cuda.empty_cache()


Even after deleting everything and calling empty_cache(), the VRAM just keeps going up and I don't understand why. This doesn't happen on CPU. If I run nvidia-smi, the memory usage increases after every batch until it crashes.

I tried:

• Calling del on everything after every batch
• Setting num_workers=0 (didn't help)
• Using .detach() before moving tensors to GPU
• Checked if the issue is in my model, but even without the model, just loading the Data already makes the memory increase

Anyone seen this before? Is there something about DataLoader and cuda() that could be causing this?

Would appreciate any ideas. I'm out of things to try

Continue reading...