请使用 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

它给了我以下警告：

=> 运行时错误：尝试在 CUDA 设备上反序列化对象，但 torch.cuda.is_available() 为 False。如果您在仅使用 CPU 的计算机上运行，请使用 torch.load 和 map_location=torch.device('cpu') 将您的存储映射到 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')

Answer 2

只是给出一个较小的答案。要解决此问题，您可以更改

load()

文件中名为

serialization.py

的函数的参数。这存储在：

./site-package/torch/serialization.py

写：

def load(f, map_location='cpu', pickle_module=pickle, **pickle_load_args):

而不是：

def load(f, map_location=None, pickle_module=pickle, **pickle_load_args):

希望有帮助。