Should be a way to unpickle an object with a torch cuda tensor on a CPU-only machine when using plain "pickle"

[petered]

🚀 Feature

The ability to pickle.load a Python object containing a torch cuda tensor on a CPU only machine.

Motivation

Currently, trying to do this gives 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='cpu' to map your storages to the CPU. Even though you are loading with pickle.load, not torch.load.

When the "loading" code is pytorch agnostic (exists in a repo that does not use pytorch), you can't just change a pickle.load(f) into a torch.load(f, map_location='cpu'). This may be the case when the saved data takes a particular structure and loading/unloading is handled by some code that does not depend on pytorch.

Pitch

A context manager could take care of this:

with torch.loading_context(map_location='cpu'):
    obj = pickle.load(f)  # In my case this call is buried deeper in torch-agnostic code

[valeriu-balaban]

A handy feature when you do the training on a machine with a GPU but then do the plotting on your laptop which does not have a GPU.

[JonathanCMitchell]

Has this been solved?

[andres-fr]

Is there an update here?

[avacariu]

I'm seeing this issue even when using torch.load(f, map_location='cpu'); it seems it loses the map_location somewhere along the way. In my case the pickle file is storing { "state_dict": model.state_dict() }.

Specifically loading fails on this line (since it executes torch.load() without passing the map_location that was passed to the parent call of torch.load()):

pytorch/torch/storage.py

Lines 135 to 136 in a7cf5c8

	def _load_from_bytes(b):
	return torch.load(io.BytesIO(b))

[mmcdermott]

For errors caused by pickle.load, you can workaround this with a custom Unpickler:

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

It may be the case that in some applications it is a different module/name causing issues, but for me it was torch.storage/_load_from_bytes so that's all I corrected.

[janblumenkamp]

@mmcdermott Are you aware of a similar workaround for loads? pickle.Unpickler seems to only work for load.

[mmcdermott]

@janblumenkamp I'm not, sorry.

[andreynwix]

why not implement a simple solution in storage.py:

if torch.cuda.is_available():
    return torch.load(io.BytesIO(b))
else:
    return torch.load(io.BytesIO(b), map_location=torch.device('cpu'))

this will remove any need to use custom loaders.
thoughts?