I am training a model on the gpu and after each epoch I would like to store the best models on the cpu. Does model.to(‘cpu’) create a copy on the cpu and does it keep the original model on the gpu?
At the moment, I have this code:
best_model = copy.deepcopy(model)
best_model = best_model.to(‘cpu’)
The problem is that this code makes a copy first on the gpu, and then it transfers the copy to the cpu. I would like to avoid making a copy on the gpu, since I don’t have enough gpu memory.
I think the best way is to save the model using torch.save(‘modelname.pt’). You can now load the saved model for inference. To run inference on CPU, specify map_location=‘cpu’ on torch.load()
Thank you for the reply. The solution that you propose consists of saving the best model to the disk. Do you know an alternative solution that doesn’t save the model to the disk. I just need to make a copy on the cpu of a model on the gpu, without touching the model on the gpu.
I am a bit late but I found a in-memory solution to the problem and thought to share it with others. The method does not seem very fast, but I don’t know any other at the moment.
from io import BytesIO
# Save the model to memory from GPU
model_data_in_memory = BytesIO()
torch.save(original_model, model_data_in_memory)
model_data_in_memory.seek(0)
# Load the model from memory to CPU
model_in_cpu = torch.load(model_data_in_memory, map_location="cpu")
model_data_in_memory.close()
The seek is needed to seek the in-memory file to the beginning. After that just load the model to cpu assuming that the original model was on the GPU. Be sure to also close the file to free the memory so that the buffer doesn’t leave open.
Edit: I found that if you add to “pickle_protocol=-1” to the save command it speeds up about 10% (in pytorch 1.13).