Save model state in memory (cpu and gpu)

1

I would like to save and load the weights of my model several time for each epochs during training. At the moment I save and load the model from a file but it is slow.
I know that it is possible to save it inside a IOBuffer object but it is not possible on GPU. Is there a way to save and load a model from a variable on both CPU and GPU quickly ?

Thank you

2

Just assign model.state_dict() to a variable, and then call model.load_state_dict() on the variable to load the weights back. See the example below.

from torchvision.models import resnet
states =[]
model = resnet.resnet18()

states.append(model.state_dict())

model.load_state_dict(states[0])

Edit: If you want to store the weight on the GPU then:

model = resnet.resnet18().cuda()
model2 = model
...
model = model2
del model2
3

Thank you.
But it does not work, the weights of the model do not change after the loading for both solutions :

model = torch.nn.Linear(4,4)
save = model
#training here
print(net.weight.data)
model = save
print(net.weight.data)
#Same result for both print

model = torch.nn.Linear(4,4)
save = model.state_dict()
#training here
print(net.weight.data)
model.load_state_dict(save)
print(net.weight.data)
#Same result for both print