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

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.

所以我使用了以下 with 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当关卡加载。

这意味着您需要手动加载检查点并将其应用于模型。值得庆幸的是，模型可以在不加载检查点的情况下加载相同的配置，通过设置pretrained=False和检查点可以单独加载torch.hub.load_state_dict_from_url，支持map_location. 状态 dict 只需要一个小的改变，因为它是用 训练的DistributedDataParallel，它包装了模块，这样layer模型中的每个都变成了module.layer。在状态 dict 中，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)