Load trained weights from RAM instead of from disk

During my research work, I need to iterate and modify a python script frequently. In this script, I load trained weights, and it takes about 3 minutes to load those weights in RAM/GPU. 3 minutes is too long of an interruption for just waiting, so I need to switch back and forth between this task and something else.

I wish I could keep those trained weights in RAM in between runs of my script! Anybody knows how I could achieve this? I can’t be the only one in that situation!

I would need some kind of database process that would be continuously running, and that would pass data to my script. Any thoughts?

I think the simplest way to achieve this would be to use a ramdisk.
You can look this up on google but the idea is that some part of your filesystem is backed by RAM instead of an HDD. So loading from a file in this part of your filesystem would be loading from RAM directly.

That being said it sound weird that loading takes so long.
How big is your network? Is it really the torch.load call which is slow? Or something else?

1 Like

Thanks for the quick answer.

The problem was that I wasn’t using torch.load… Since I wanted to save the state_dict along with some training/validation curves, I was saving a tuple (state_dict, results) using pickle

I just tested and successfully saved a tuple like this with torch.save. I didn’t know this function was equivalent to pickle.dump ! Loading it back takes a split second.

Have a good day, sir! :slight_smile:

Ho yes, pickle doesn’t know about torch.Tensor so takes forever to save them.
torch.load is actually an extension of pickle that supports that efficiently.

Glad it works well now !

2 Likes

No way to accept an answer, mark an answer as solution, mark a question as solved?

No I don’t think so :confused: