如何使用 map_location='cpu' 由于“RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False”

Question

我试图在https://pytorch.org/hub/nvidia_deeplearningexamples_tacotron2/下载以下模型

import torch
tacotron2 = torch.hub.load('nvidia/DeepLearningExamples:torchhub', 'nvidia_tacotron2')

我收到了：

>>> import torch
>>> tacotron2 = torch.hub.load('nvidia/DeepLearningExamples:torchhub', 'nvidia_tacotron2')
Using cache found in .cache\torch\hub\nvidia_DeepLearningExamples_torchhub
...
  File "Anaconda3\envs\env3_pytorch\lib\site-packages\torch\serialization.py", line 79, in validate_cuda_device
    raise RuntimeError('Attempting to deserialize object on a CUDA '
RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location='cpu' to map your storages to the CPU.

所以我将以下内容与map_location='cpu' 一起使用，但仍然得到相同的错误。

>>> tacotron2 = torch.hub.load('nvidia/DeepLearningExamples:torchhub', 'nvidia_tacotron2', map_location='cpu')

torch.version.cuda 显示一个版本，但 torch.cuda.is_available() 为假。

>>> import torch
>>> torch.version.cuda
'9.0'
>>> torch.cuda.is_available()
False

• 当我使用了要求使用的命令错误消息时，如何解决与map_location 相关的错误？

• 如何使用 cuda 加载模型？

Answer 1

torch.hub.load 不具体支持map_location，它只是将额外的参数转发给模型的加载，因此它是否支持取决于实现。

在这种情况下不支持，加载是在NVIDIA/DeepLearningExamples:torchhub - hubconf.py 中实现的，它不会将任何map_location 传递给torch.load when the checkpoint is loaded。

这意味着您需要手动加载检查点并将其应用于模型。值得庆幸的是，通过设置pretrained=False，可以在不加载检查点的情况下加载相同的模型，并且可以使用支持map_location的torch.hub.load_state_dict_from_url单独加载检查点。状态字典只需要做一个很小的改变，因为它是用 DistributedDataParallel 训练的，它包装了模块，这样模型中的每个 layer 都变成了 module.layer。在状态字典中，module. 前缀需要被剥离。

tacotron2 = torch.hub.load('nvidia/DeepLearningExamples:torchhub', 'nvidia_tacotron2', pretrained=False)

checkpoint = torch.hub.load_state_dict_from_url('https://api.ngc.nvidia.com/v2/models/nvidia/tacotron2pyt_fp32/versions/1/files/nvidia_tacotron2pyt_fp32_20190306.pth', map_location="cpu")

# Unwrap the DistributedDataParallel module
# module.layer -> layer
state_dict = {key.replace("module.", ""): value for key, value in checkpoint["state_dict"].items()}

# Apply the state dict to the model
tacotron2.load_state_dict(state_dict)