Hi guys.
I am working on a Python module that would allow me to train multiple builtin torchvision models in sequence (one after another) keeping track of training results. Sometimes, I just need a rough idea of what network would perform better for my problem, so that I can more carefully play with it later.
The problem:
When it comes to training I have a list of models (resnets, vggs, alexnet, inception, densenet and squeezenet), a basic FOR loop to train one network after another and a dictionary where I track training results of each model. I later visualize them. I made sure everything works fine, first training sessions went pretty.
But there is a problem, now very quickly models stop to train (accuracy around 45-55%, random guesses for 2 classes), first 2 models in the list seem to train fine though. On top of that, on around 4th network in my list (VGG19) I get an error - out of memory, even though when I train only VGG19 alone everything works fine. Clearly, models share some resource which gets filled at some point and training becomes impossible.
The question: could anyone please help me identify those memory leaks and fix the issue?   My guesses:

All models get moved to GPU. Should I remove it from there at the end of each training?
Optimizer, loss function and scheduler get initialized for each new model to train, so I think there should be no problem. Correct?
During training batches and labels get moved to GPU as well. Should I remove them at the end, so that GPU memory is available for the next training session or it gets cleaned automatically?

In case anyone wishes to have a look at the code -  https://github.com/EvgeniiTitov/Neural-Net-Trainer 6
Would love to hear adequate criticism from you people. The trainer class can be found in wrappers/ line 291 - GroupTrainer. It gets called from the main.py
Thank you very much guys!
Regards,
Eugene