FIX Unpickling without using torch.load

[BenjaminBossan]

Resolves #1090, #1091.

This one got more complicated than I initially thought. So here it goes:

Description

PyTorch plans to make the switch to weights_only=True for torch.load. We already partly dealt with that in #1064 when it comes to save_params/load_params. However, we still had a gap. Namely, when using pickle directly, i.e. when going through __getstate__ and __setstate__, we are still using torch.load and torch.save without handling weights_only. This will cause trouble in the future when the default is switched. But it's also annoying right now, because users will get the FutureWarning about weights_only, even if they correctly pass torch_load_kwargs (see #1090).

The reason why we use torch.save/torch.load for pickle is that those functions are basically extended pickle functions that have the benefit of supporting the map_location argument to handle the device of torch tensors, which we don't have for pickle. The map_location argument is important, e.g. when saving a net that uses CUDA and loading it on a machine without CUDA, we would otherwise run into an error.

However, with the move to weights_only=True, these torch.save/torch.load will become reduced pickle functions, as they will only support a small subset of objects by default. Therefore, we wouldn't be able to rely on torch.save/torch.load for pickling the whole skorch object.

Solution

(thanks ChatGPT o1 for helping with this)

In this PR, we move to using plain pickle for this. However, now we run into the issue of how to handle the map_location. The solution I ended up with is now to intercept torch's _load_from_bytes using a custom Unpickler, and to specifically use torch.load there. That way, we can pass the map_location and other torch_load_kwargs. The remaining unpickling process just works as normal.

Yes, this is a private function, so we cannot be sure if it'll work indefinitely, If there is a better suggestion, I'm open to it. However, the function has existed for 7 years, so it's not very likely that it will change anytime soon:

https://github.com/pytorch/pytorch/blame/0674ab7e33c3f627ca6781ce98468ec1dd4743a5/torch/storage.py#L525

A drawback of the solution is that we cannot just load old skorch nets that were saved with torch.save using pickle.load. This is because torch uses custom persistent_load functions. When trying to load with pickle, we thus get:

_pickle.UnpicklingError: A load persistent id instruction was encountered, but no persistent_load function was specified.

Therefore, I had to keep torch.load as a fallback to avoid backwards incompatibility. The bad news is that the initial problem persists, namely that even when passing torch_load_kwargs, users get the FutureWarning about weights_only. The good news is that users can just re-save their net with the new skorch version and from then on they won't see the warning again.

Note that I didn't add a specific test for this problem of loading nets from before the change, because test_pickle_load, which uses a checked in pickled net, already covers this.

Other considered solutions

1. Why not continue using torch.save/torch.load and just pass the torch_load_kwargs argument to it? This is unfortunately not that easy. When switching to weights_only=True, torch will refuse to load any custom objects, e.g. class MyModule. There is a way to prevent that, namely via torch.serialization.add_safe_globals, but it is a ton of work to add all required objects there, as even builtin Python types are mostly not supported.
2. We cannot use with torch.device(map_location):, as this is not honored during unpickling.
3. During __getstate__, we could recursively go through the state, pop all torch tensors, and replace them with, say, numpy arrays and additional meta data like the device, then use this info to restore those objects during __setstate__. Even though this looks like a cleaner solution, it is much more complex and therefore, I'd argue, more error prone.
4. Edit: Don't do anything and just live with the warning: This will work -- until PyTorch switches the default. Therefore, we had to tackle this sooner or later.

Notes

While working on this, I thought that we could most likely remove the cuda_dependent_attributes_ (which contains the net.module_, net.optimizer_, etc.). Their purpose was to call torch.load on these attributes specifically, but with the new Unpickler, it should also work without this. However, I kept the attribute for now, mainly for these reasons:

1. I didn't want to change more than necessary, as these changes are delicate and I don't to break any existing skorch code or pickle files.
2. The attribute itself is public, so in theory, users may rely on its existence (not sure if in practice). We would thus have to keep most of the code related to this attribute anyway.

But LMK if you think we should deprecate and eventually remove this attribute.

[githubnemo]

LGTM, seems to be a good complexity/impact trade-off.