python - 请将 torch.load 与 map_location=torch.device ('cpu' 一起使用)

Question

我正在运行 Python 程序，但我没有 GPU，我该怎么做才能让 Python 使用 CPU 而不是 GPU？

$ python extract_feature.py --data mnist --net checkpoint_4.pth.tar --features pretrained

它给了我以下警告:

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

照片是我的 Python 项目的结构:

Answer 1

我遇到了类似的错误。然后通过尝试以下解决方法问题得到解决。 (如果您的模型是 .pth 或 .h5 格式。)

MODEL_PATH = 'Somemodelname.pth'
model.load_state_dict(torch.load(MODEL_PATH, 
map_location=torch.device('cpu')))

如果您希望在您的机器中使用特定的 GPU。然后，

map_location = torch.device('cuda:device_id')