[QUESTION] Infer on CPU while model is trained in GPU

[gsamaras]

I had trained an N-Beats model in GPU and now I want to use to do predictions in CPU.

from darts.models.forecasting.nbeats import NBEATSModel

model_nbeats = NBEATSModel(input_chunk_length=3,  output_chunk_length=1, generic_architecture=True)
model_nbeats.save_model("nbeats_model.pth.tar") # Save on GPU

Then on a CPU machine I do:

model_nbeats = NBEATSModel.load_model("nbeats_model.pth.tar")

and get this error:

RuntimeError: Attempting to deserialize object on a CUDA device

---

My attempt was to follow this SO question, like:

model_nbeats = load_model("nbeats_model.pth.tar", map_location=torch.device('cpu'))

but this method doesn't accept any such parameter.

If I then use torch.load(), like:

import torch
model_nbeats = torch.load("nbeats_model.pth.tar", map_location=torch.device('cpu'))

it works but then the predict method doesn't work:

current_pred_series = model_nbeats.predict( series = test_series, n = 1)

because it's a Torch model and not a Darts one.

Can you please help?

---

I cannot use the Trainer parameter in predict(), because I am using Darts 0.16.1 and cannot change at this time, which has this signature:

(n: int, series: Union[darts.timeseries.TimeSeries, Sequence[darts.timeseries.TimeSeries], NoneType] = None, past_covariates: Union[darts.timeseries.TimeSeries, Sequence[darts.timeseries.TimeSeries], NoneType] = None, future_covariates: Union[darts.timeseries.TimeSeries, Sequence[darts.timeseries.TimeSeries], NoneType] = None, batch_size: Union[int, NoneType] = None, verbose: bool = False, n_jobs: int = 1, roll_size: Union[int, NoneType] = None, num_samples: int = 1, num_loader_workers: int = 0) -> Union[darts.timeseries.TimeSeries, Sequence[darts.timeseries.TimeSeries]]

[dennisbader]

Hey @gsamaras, have you tried once with checkpointing and then NBEATSModel.load_from_checkpoint()?

We switched to PyTorch Lighting a couple of versions ago which changed the save/load logic.

If this doesn't work and you can't upgrade your Darts version, then you could theoretically change the TorchForecastingModel.load_model() function signature in the source code (keep in mind all the downsides of doing this) to pass the device mapping to the internal torch.load() call

[gsamaras]

Hi @dennisbader, so you mean, with Darts 0.16.1, in the GPU machine use NBEATSModel.save_checkpoints() (ref), and then in the CPU machine use NBEATSModel.load_from_checkpoint()?

Yes currently upgrading is not an option for my project-related-internal-reasons. :/

[dennisbader]

You should activate checkpointing at model creation on the GPU machine (see model docs here):

model = NBEATSModel(..., model_name="my_model", work_dir=my_work_dir, save_checkpoints=True)

And then on the CPU machine load from checkpoint see here

model = NBEATSModel.load_from_checkpoint(model_name="my_model", work_dir=my_work_dir, best=False)

(I can't say for certain that it will work though)

[gsamaras]

It didn't. On load I got:

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.

If you have other ideas please let me know. Otherwise, when time permits in the future, I will upgrade Darts and revisit this.

[dennisbader]

Then you would have to adapt the code:

• fork the repo
• branch out of master
• reset the head to where it was in your specific darts version
• adapt the loading methods to pass the device mapping to the internal torch loading
• install darts locally from this branch

[gsamaras]

Thank you, I'll see if I manage to modify darts or if I'll be able to update Darts later.