fix: torch_forecasting_model load_weights with float16 and float32

[Hsinfu]

Fix bug for TorchForecastingModel calling load_weights with float16 and float32. Fixes #1987.

Bug details

• Train the TorchForecastingModel with float32 following the instruction and .save it to disk.
• Create a new model instance and call .load_weights to restore it.
• The mock_train_sample in .load_weights_from_checkpoint built from np.zero without specifying the dtype, so it sticks to the default float64, and causes the self.train_sample stays float64, too.

[dennisbader]

Thanks @Hsinfu for contributing and this fix 🚀

I had 2 minor suggestions regarding adding a unit test and making the dtype mapping more robust against future changes.

[dennisbader]

Can you try adding a unit test for this?

The best place for this would be in darts/tests/models/forecasting/test_torch_forecasting_model.
Maybe something like test_loading_weights_with_different_dtype

[Hsinfu]

test_load_weights_with_float32_dtype

[dennisbader]

I'd suggest creating a mapping between the torch and numpy dtypes at the top of this file. With this we can avoid/catch some breaking changes (e.g. if we ever support other dtypes such as float16 and forget about this line).

TORCH_NP_DTYPES = {
    torch.float32: np.float32,
    torch.float64: np.float64,
}

...

mock_train_sample = [
    np.zeros(sample_shape, dtype=TORCH_NP_DTYPES[ckpt["model_dtype"]]) if sample_shape else None
    for sample_shape in ckpt["train_sample_shape"]
]

[Hsinfu]

also add float16

TORCH_NP_DTYPES = {
    torch.float16: np.float16,
    torch.float32: np.float32,
    torch.float64: np.float64,
}

[madtoinou]

Just for the record, this PR fixed one of the point mentioned in #1987.

[codecov-commenter]

Codecov Report

Attention: 2 lines in your changes are missing coverage. Please review.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Files	Coverage Δ
darts/dataprocessing/encoders/encoders.py	97.68% <100.00%> (+0.07%)	⬆️
darts/models/forecasting/arima.py	93.75% <ø> (ø)
darts/models/forecasting/auto_arima.py	96.77% <ø> (ø)
darts/models/forecasting/block_rnn_model.py	98.27% <ø> (-0.03%)	⬇️
darts/models/forecasting/catboost_model.py	98.41% <ø> (ø)
darts/models/forecasting/croston.py	90.69% <ø> (ø)
darts/models/forecasting/dlinear.py	99.12% <ø> (ø)
darts/models/forecasting/kalman_forecaster.py	97.36% <ø> (ø)
darts/models/forecasting/lgbm.py	100.00% <ø> (ø)
...arts/models/forecasting/linear_regression_model.py	94.44% <ø> (ø)
... and 18 more

... and 2 files with indirect coverage changes

📢 Thoughts on this report? Let us know!.

[dennisbader]

Hi @Hsinfu , it seems like your branch is protected for us to adapt it.

Can you merge the current master head into your branch and add the following line to our CHANGELOG.md Fixed section?

- Fixed a bug when loading a `TorchForecastingModel` that was trained with a precision other than `float64`. [#2046](https://github.com/unit8co/darts/pull/2046) by [Freddie Hsin-Fu Huang](https://github.com/Hsinfu).

[Hsinfu]

• Fixed a bug when loading a TorchForecastingModel that was trained with a precision other than float64. #2046 by Freddie Hsin-Fu Huang.

I just cloned the current master head, and rebase my branch on it. Also add the information in CHANGELOG.md.
@dennisbader See if it works now. Thank you~

[dennisbader]

Looks great, thanks a lot for this @Hsinfu 🚀 💯