I am dealing with variable sized input, and most instances are too large to fit in memory, even using multiple GPUs. I therefore use one network to compute a reduced representation (embedding) which I then forward to the prediction head (basically a form of gradient checkpointing).
There are surely many ways to implement this.
I am particularly surprised with a memory leak I encountered in the following approach:
self.network.eval()
with torch.no_grad():
instances = defaultdict(list)
for batch in dataloader:
img = batch["img"].to(self.device)
batch["embedding"] = self.network(img).detach().cpu()
for key in patchbatch:
instances[key].append(batch[key])
I believe that tensor.to.self(device) is not an in-place operation and results in a copy. So when I later collect batch[“img”] and batch[“embedding”] for a later backwards pass, they do not live on the GPU.
At which place am I blocking the GPUs memory from being freed?