当我将 pickle 文件加载到 CPU 时遇到问题。我在网上搜索了一下,他们说我需要添加map_location参数。但是我添加这个参数后,问题依然存在。
代码如下:
torch.__version__
torch.load('featurs.pkl',map_location='cpu')
>>>
'1.0.1.post2'
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=torch.device('cpu')
但即使这样也不起作用。一种解决方法是使用 pickle
库并实现自定义 unpickler。
import pickle
import io
class CPU_Unpickler(pickle.Unpickler):
def find_class(self, module, name):
if module == 'torch.storage' and name == '_load_from_bytes':
return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
else:
return super().find_class(module, name)
#contents = pickle.load(f) becomes...
contents = CPU_Unpickler(f).load()
来源:Github
试试这个:
torch.load('featurs.pkl',map_location=torch.device('cpu'))