Runtimeerror: Is there a way to torch.load(model.pt) on CPU?

user_12062371 0
Viewed 1078

For the usage of my GPU-trained model I want to run in on my CPU. If I understand it right, there is already a solution for that, which is to save and load a model parameters (or state_dict). The loading process (which is what I am interested here) would then go something like this:

device = torch.device('cpu')
state_dict = torch.load(self.model_state_dict_path, map_location=self.device)

model = ModelClass(*model_params*)
model.load_state_dict(state_dict)
model.to(device)

, which seems to work fine.

I was wondering, though, if it was also possible to achieve the same by saving/loading the entire model and not its parameters. Trying to load the GPU-trained model on the CPU, I did the following:

    def __init__(self, entire_model_path):
        self.model_path = entire_model_path
        self.__device = device('cpu')
        self.__model = load(model_path, map_location=self.__device)

, which throws the following error:

Traceback (most recent call last):
  File "predictor.py", line 138, in <module>
    p.predict_masks_from_directory(config.image_path, config.mask_path)
  File "predictor.py", line 25, in predict_masks_from_directory
    pred = self.predict_mask_from_imagepath(imagepath)
  File "predictor.py", line 37, in predict_mask_from_imagepath
    pred_patches.append(np.moveaxis(self.predict_segmentation_from_image(img)[0], 0, -1))
  File "predictor.py", line 53, in predict_segmentation_from_image
    segmentation_result = sigmoid(self.__unet(image))
  File "C:\Users\bboche\Anaconda3\envs\UNet\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "C:\Users\bboche\Anaconda3\envs\UNet\lib\site-packages\torch\nn\parallel\data_parallel.py", line 155, in forward
    "them on device: {}".format(self.src_device_obj, t.device))
RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu

Is it just not possible to load an entire model in that way, since it has saved some data that is GPU-specific, maybe?

Also: is it even desirable to load an entire model instead of its state dictionary?

pytorch python
1 Answers

Not sure whether your question comes because of habits using TensorFlow SavedModel for instance, but PyTorch is usually about saving/loading parameters' values.

Moving a model to a device is effectively moving all its parameters (values & gradients) to the target device. So, apart from if it's really time-consuming for you, the best option is usually to:

  • instantiate your model on CPU
  • load your checkpoint on CPU
  • load the param values from the checkpoint onto your model
  • moving your model to the target device
import torch
model = ...
state_dict = torch.load("path/to/your/checkpoint.pt", map_location="cpu")
model.load_state_dict(state_dict)
model = model.to(device=...)

Now if you want to save an entire model (architecture + its weights), you can do it (cf. https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model) as follows:

import torch
model = ...
# Save the architecture and the weights
torch.save(model, "path/to/your/checkpoint.pt")

# Load it
restored_model = torch.load("path/to/your/checkpoint.pt")
user_109746650
Related