使用 CPU 加载 pickle 保存的 GPU 张量？

Question

我在 GPU 上使用 pickle 保存了 Bert 的最后一个隐藏层，用于后续流程。

# output is the last hidden layer of bert, transformed on GPU
with open(filename, 'wb') as f:
    pk.dump(output, f)

是否可以将其加载到我没有 GPU 的笔记本电脑上？我尝试了以下代码，但都失败了。

# 1st try
with open(filename, 'rb') as f:
    torch.load(f, map_location='cpu')

# 2nd
torch.load(filename, map_location=torch.device('cpu'))

都出现以下错误

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.

可以将文件加载到我的笔记本电脑上吗？

Answer 1

如果您使用 pytorch，您可以通过保存模型的

state_dict

而不是模型本身来省去一些麻烦。

state_dict

是一个有序字典，用于存储神经网络的权重。

保存程序：

import torch
model = MyFabulousPytorchModel()
torch.save(model.state_dict(), "best_model.pt")

加载它需要您先初始化模型：

import torch
device = 'cuda' if torch.cuda.is_available() else 'gpu' 
model = MyFabulousPytorchModel()
model.load_state_dict(torch.load(PATH_TO_MODEL))
model.device(device)

保存

state_dict

而不是直接保存对象有很多优点。其中之一与您的问题有关：将模型移植到不同的环境并不像您希望的那样轻松。另一个优点是，保存检查点要容易得多，让您可以恢复训练，就好像训练从未停止过一样。您所要做的就是保存优化器的状态和损失：

保存检查点：

# somewhere in your training loop:
opt.zero_grad()
pred = model(x)
loss = loss_func(pred, target)

torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "loss":loss}, "checkpoing.pt")

我强烈建议您查看文档以获取有关如何使用 pytorch 保存和加载模型的更多信息。如果你了解其内部运作原理，这是一个相当顺利的过程。 https://pytorch.org/tutorials/beginner/ saving_loading_models.html# saving-loading-model-for-inference

希望有帮助=）

编辑：

更直接地，为了解决您的问题，我推荐以下内容

1- 在您用来训练模型的计算机上：

import torch
model = torch.load("PATH_TO_MODEL")
torch.save(model.state_dict(), "PATH.pt")

2- 在另一台计算机上：

import torch
from FILE_WHERE_THE_MODEL_CLASS_IS_DEFINED import Model

model = Model() # initialize one instance of the model)
model.load_state_dict(torch.load("PATH.pt")

Answer 2

我认为您可以通过在调用之前包含此代码来解决此问题

torch.load

：

torch.serialization.register_package(0, lambda x: x.device.type, lambda x, _: x.cpu())

这会注册一个自定义反序列化器，将存储对象转换为 CPU。

尽管如此，我认为因为您已经调用

pk.dump

来保存对象，所以您还应该调用

pk.load

来加载它（而不是

torch.load

）。