Pytorch/cuda：CPU 错误和 map_location

如何解决Pytorch/cuda：CPU 错误和 map_location

我编写此代码来下载我的模型：

args = parser.parse_args()

use_cuda = torch.cuda.is_available()

state_dict = torch.load(args.model)
model = Net()
model.load_state_dict(state_dict)
model.eval()

if use_cuda:
    print('Using GPU')
    model.cuda()
else:
    print('Using cpu')

但我的终端返回以下错误 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=torch.device('cpu') to map your storages to the cpu.

所以我试着在没有真正理解太多的情况下写作：

args = parser.parse_args()

map_location=torch.device('cpu')
state_dict = torch.load(args.model)
model = Net()
model.load_state_dict(state_dict)
model.eval()

但我还是犯了同样的错误。你看我如何纠正它？ （实际上我想用我的 cpu 加载我的模型）。

解决方法

我假设您将模型保存在带有 GPU 的计算机上，现在正在一台没有 GPU 的计算机上加载它，或者您可能由于某种原因 GPU 不可用。另外，哪一行导致了错误？

参数map_location需要设置在torch.load里面。像这样：

state_dict = torch.load(args.model,map_location='cpu')

或

map_location=torch.device('cpu')
state_dict = torch.load(args.model,map_location=map_location)

请注意，您需要将 map_location 变量发送到 torch.load 函数。