Load pickle saved GPU tensor with CPU?

Question

I save the last hidden layer of Bert for my following process using pickle on GPU.

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

Is it possible to load it on my person laptop without GPU? I tried following code, but all failed.

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

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

All get the following error

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.

Is it possible to load the file on my laptop?

Woodenhead

Answer 1

If you using pytorch, you can save yourself some headache by saving the state_dict of the model instead of the model itself. The state_dict is an ordered dictionary that stores the weights of your neural network.

The saving routine:

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

Loading it requires you to initialize the model first:

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)

There are many advantages of saving the state_dict instead of the object directly. One of them has to do with your issue: porting your model to a different environment isn't as painless as you were hoping for. Another advantage is that it is a lot easier to save checkpoints that allow you to resume training as if training had never stopped. All you have to do is save the state of the optimizer and the loss:

Saving checkpoint:

# 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")

I highly recommend checking the documentation for further information on how to save and load models using pytorch. It is quite a smooth process if you understand its inner workings. https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-for-inference

Hope that helps =)

Edit:

More directly, to solve your problem I recomend the following

1- On the computer you used to train the model:

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

2- On the other computer:

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")

Mandrill

I pickled a cuda tensor which I want to load on CPU. Using the same 2 methods given by OP doesn't work. What can be done being not a state_dict()? I saved using pickle.dump(f, 'some_cuda_tensor.pkl') – Dryden

OP explicitly mentions a tensor, not a model. – Hoxsie

@Gaston, the question regards loading a model. However, you can use `torch.save(whatever_tensor, "tensor.pt") to saving a tensor as well. – Mandrill

Answer 2

I think you may be able to solve this by including this code before you call torch.load:

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

This registers a custom deserializer that converts the storage object to CPU.

Although, I think because you have called pk.dump to save the save the object, you should also call pk.load to load it (rather than torch.load).

Damle