如何反序列化含未转CPU的PyTorch模型的Pickle字典？

方法一：使用torch.load指定map_location #

虽然你用Pickle保存的字典，但PyTorch张量（包括模型参数）的序列化格式与torch.load兼容。直接用torch.load加载并指定map_location为CPU，就能自动将所有CUDA张量映射到CPU：

import torch
import pickle

with open("你的文件路径.pkl", "rb") as f:
    loaded_dict = torch.load(f, map_location=torch.device('cpu'))

加载完成后，字典内的模型和所有张量都会自动切换到CPU设备，可正常使用。

方法二：自定义Pickle反序列化钩子 #

如果必须使用pickle.load，可以注册一个对象钩子，在反序列化时自动将CUDA张量转成CPU张量：

import pickle
import torch

def convert_cuda_to_cpu(obj):
    if isinstance(obj, torch.Tensor) and obj.is_cuda:
        return obj.cpu()
    return obj

with open("你的文件路径.pkl", "rb") as f:
    loaded_dict = pickle.load(f, object_hook=convert_cuda_to_cpu)

这个钩子会遍历反序列化的每个对象，仅处理CUDA张量，其余内容保持原样。加载后若需使用模型，可手动执行loaded_dict["你的模型键"].cpu()确保模型完全在CPU上。

关键提示 #

• 两种方法都能解决CPU环境下的反序列化问题，优先推荐方法一，因为torch.load对PyTorch对象的兼容性更好。
• 若字典内包含其他CUDA相关非张量对象，需针对性扩展钩子函数，但绝大多数场景下上述方法足够覆盖需求。

内容的提问来源于stack exchange，提问作者user202542