运行时错误：有没有办法在 CPU 上执行 torch.load(model.pt)？

Question

为了使用我的 GPU 训练模型，我想在我的 CPU 上运行。 如果我理解正确，那么已经有一个解决方案，即保存和加载模型参数（或state_dict ）。 加载过程（这是我在这里感兴趣的）然后会是这样的：

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)

，这似乎工作正常。

不过，我想知道是否也可以通过保存/加载整个模型而不是其参数来实现相同的效果。 尝试在 CPU 上加载 GPU 训练模型，我执行了以下操作：

    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)

，这会引发以下错误：

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

是否不可能以这种方式加载整个模型，因为它可能保存了一些特定于 GPU 的数据？

另外：是否需要加载整个模型而不是其状态字典？

Answer 1

不确定您的问题是否是因为使用 TensorFlow SavedModel 的习惯，但 PyTorch 通常是关于保存/加载参数的值。

将模型移动到设备实际上是将其所有参数（值和梯度）移动到目标设备。 因此，除了对您来说真的很耗时之外，最好的选择通常是：

• 在 CPU 上实例化您的模型
• 在 CPU 上加载检查点
• 将检查点中的参数值加载到您的模型上
• 将模型移动到目标设备

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=...)

现在，如果您想保存整个模型（架构 + 其权重），您可以这样做（参见https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model ），如下所示：

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")