How to delete eveything except the model from the GPU without restarting the ipython kernel?

1

I’m trying to experiment with different model-architectures in a single jupyter notebook.

At the end of the notebook, I want to run all the n-models I trained above on the test-data.

But my GPU doesn’t have enough memory to train n-different models at once. I need to free up some space before training the next model. But I also need to preserve my n-models on the GPU so that I can use them all later for inference/test.

This is something I’m going for -

train_dataloader = DataLoader(...)
test_dataloader = DataLoader(...)


# model-1
class Model(nn.Module):
    def __init__(self):
        bla = nn.Conv2d()
    def forward(self, x):
        x = bla(x)
        return x


# train model_1
for i in trainloader:
    # train model_1
    pred  = model_1(x)

# save model_1 params and free up the GPU

# train model_2
# save model_2 params and free up the GPU


# train model_3
# save model_3 params and free up the GPU


# finally test all three models together

How do free up space on the GPU without deleting the model(on the gpu) and without restarting the kernel?

I came across torch.cuda.empty_cache(). Can I safely use it in this case?

2

theairbend3r,

you could try saving the trained model & optimiser state dictionaries and then load as needed.

https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model

3

In this way, the GPU would still have all the tensors/models. I would need to restart my kernel to free up the memory, right?

4

if you save the tensordata set as well using torch.save(tensordatasetnetname, ‘filename.pt’) i guess so

5

Hmm. I want to avoid restarting the kernel. I’d like to train all the n-models in one go. Is that possible? torch.cuda.empty_cache() seems to work. I basically call it after training each model. But I’m not sure what it’s clearing.