gallantlab / himalaya

Multiple-target linear models - CPU/GPU
https://gallantlab.github.io/himalaya
BSD 3-Clause "New" or "Revised" License
78 stars 13 forks source link

Convert model from GPU to CPU #52

Closed tommybotch closed 7 months ago

tommybotch commented 7 months ago

Hello!

I am attempting to reload a model trained on a GPU within CPU memory using joblib, however I receive 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.

Generally, torch models can be transferred to CPU memory using model.to_cpu(), but as the model is wrapped in an sklearn pipeline it doesn't seem this works. Is there a global function that can map the kernelizer/weights within the pipeline into CPU memory for saving?

Thanks in advance!

mvdoc commented 7 months ago

Hi Tommy,

I don't think we ever tested dumping the model with joblib. These models are much simpler than DNNs, so usually saving the individual model parameters (mostly the dual weights and the optimal alphas) is enough to re-initialize the model by manually setting the parameters in the new model instance. Implementing something like model.to_cpu() would definitely be more straightforward, but I'm afraid at the moment we don't have the bandwidth to do that.

Alternatively, I wonder if converting all the model parameters to CPU with backend.to_numpy(param) before dumping it with joblib would solve this issue.

EDIT: to clarify, you could do what I'm suggesting even if the models are within a pipeline. You can still access the individual steps of a pipeline and modify the parameters. For example, assuming the model is the last step, you should be able to access and move the dual coef to cpu with pipeline[-1].dual_coef_ = backend.to_numpy(pipeline[-1].dual_coef_).

tommybotch commented 7 months ago

Thanks for the quick response, Matteo! Awesome -- no worries, this is super helpful. I'll play around with this and will try to create some equivalent to the torch style conversion.