加载从GPU获取的pickle文件到CPU

Question

当我将 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.

我知道这是因为设备不同，但我使用了错误消息中的指令，所以我不知道下一步如何解决。

提前致谢！

Answer 1

错误消息建议使用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

Answer 2

尝试这个：

    torch.load('featurs.pkl',map_location=torch.device('cpu'))