From 51ebbcfac8ada704939c5fa84150d6003d6a871e Mon Sep 17 00:00:00 2001 From: madtoinou Date: Mon, 24 Feb 2025 19:19:07 +0100 Subject: [PATCH 01/17] feat: adding tests for the optional dependencies --- .github/workflows/merge.yml | 2 +- darts/tests/conftest.py | 9 ++ darts/tests/optional_deps/test_onnx.py | 129 +++++++++++++++++++++++++ requirements/optional.txt | 4 + 4 files changed, 143 insertions(+), 1 deletion(-) create mode 100644 darts/tests/optional_deps/test_onnx.py create mode 100644 requirements/optional.txt diff --git a/.github/workflows/merge.yml b/.github/workflows/merge.yml index b74cd0a26f..c4a6e5ce86 100644 --- a/.github/workflows/merge.yml +++ b/.github/workflows/merge.yml @@ -54,7 +54,7 @@ jobs: elif [ "${{ matrix.flavour }}" == "torch" ]; then pip install -r requirements/core.txt -r requirements/torch.txt -r requirements/dev.txt elif [ "${{ matrix.flavour }}" == "all" ]; then - pip install -r requirements/core.txt -r requirements/torch.txt -r requirements/notorch.txt -r requirements/dev.txt + pip install -r requirements/core.txt -r requirements/torch.txt -r requirements/notorch.txt -r requirements/optional.txt -r requirements/dev.txt fi - name: "Install libomp (for LightGBM)" diff --git a/darts/tests/conftest.py b/darts/tests/conftest.py index 90bf29e20b..29059d0bf9 100644 --- a/darts/tests/conftest.py +++ b/darts/tests/conftest.py @@ -17,6 +17,15 @@ logger.warning("Torch not installed - Some tests will be skipped.") TORCH_AVAILABLE = False +try: + import onnx # noqa: F401 + import onnxruntime # noqa: F401 + + ONNX_AVAILABLE = True +except ImportError: + logger.warning("Onnx not installed - Some tests will be skipped.") + ONNX_AVAILABLE = False + tfm_kwargs = { "pl_trainer_kwargs": { "accelerator": "cpu", diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py new file mode 100644 index 0000000000..3217a611e2 --- /dev/null +++ b/darts/tests/optional_deps/test_onnx.py @@ -0,0 +1,129 @@ +from typing import Optional + +import numpy as np +import pytest + +import darts.utils.timeseries_generation as tg +from darts import TimeSeries +from darts.tests.conftest import ONNX_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs + +if not (TORCH_AVAILABLE and ONNX_AVAILABLE): + pytest.skip( + f"Torch or Onnx not available. {__name__} tests will be skipped.", + allow_module_level=True, + ) +import onnx +import onnxruntime as ort +import torch + +from darts.models import ( + BlockRNNModel, + NLinearModel, + TFTModel, + TiDEModel, +) + +# from darts.models.components.layer_norm_variants import RINorm + + +class TestOnnx: + ts_tgt = tg.linear_timeseries(start_value=0, end_value=100, length=100) + ts_pc = tg.constant_timeseries(value=123.4, length=100) + ts_fc = tg.sine_timeseries(length=100) + + @pytest.mark.parametrize( + "model_cls", + [ + BlockRNNModel, + NLinearModel, + TFTModel, + TiDEModel, + ], + ) + def test_onnx_save_load(self, tmpdir_fn, model_cls): + model = model_cls( + input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs + ) + onnx_filename = f"test_{model.name}" + model.fit() + # native inference + pred = model.predict(2) + + # model export + model.to_onnx(onnx_filename) + + # onnx model verification + onnx_model = onnx.load(onnx_filename) + onnx.checker.check_model(onnx_model) + + # manual feature extraction from the series + past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( + model=model, + series=self.ts_tgt, + past_covariates=self.ts_pc, + future_covariates=self.ts_fc, + ) + + # onnx model loading and inference + onnx_pred = self._helper_onnx_inference( + onnx_filename, past_feats, future_feats, static_feats + ) + + # check that the predictions are similar + torch.testing.assert_close(onnx_pred, pred) + + def _helper_prepare_onnx_inputs( + model, + series: TimeSeries, + past_covariates: Optional[TimeSeries] = None, + future_covariates: Optional[TimeSeries] = None, + ) -> tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: + """Helper function to slice and concatenate the input features""" + past_feats, future_feats, static_feats = None, None, None + # get input & output windows + past_start = series.end_time() - (model.input_chunk_length - 1) * series.freq + past_end = series.end_time() + future_start = past_end + 1 * series.freq + future_end = past_end + model.output_chunk_length * series.freq + # extract all historic and future features from target, past and future covariates + past_feats = series[past_start:past_end].values() + if past_covariates and model.uses_past_covariates: + # extract past covariates + past_feats = np.concatenate( + [past_feats, past_covariates[past_start:past_end].values()], axis=1 + ) + if future_covariates and model.uses_future_covariates: + # extract past part of future covariates + past_feats = np.concatenate( + [past_feats, future_covariates[past_start:past_end].values()], axis=1 + ) + # extract future part of future covariates + future_feats = future_covariates[future_start:future_end].values() + # add batch dimension -> (batch, n time steps, n components) + past_feats = np.expand_dims(past_feats, axis=0).astype(series.dtype) + future_feats = np.expand_dims(future_feats, axis=0).astype(series.dtype) + # extract static covariates + if series.has_static_covariates and model.uses_static_covariates: + static_feats = np.expand_dims( + series.static_covariates_values(), axis=0 + ).astype(series.dtype) + return past_feats, future_feats, static_feats + + def _helper_onnx_inference( + self, + onnx_filename: str, + past_feats: torch.Tensor, + future_feats: torch.Tensor, + static_feats: torch.Tensor, + ): + ort_session = ort.InferenceSession(onnx_filename) + # extract only the features expected by the model + ort_inputs = {} + for name, arr in zip( + ["x_past", "x_future", "x_static"], [past_feats, future_feats, static_feats] + ): + if name in [inp.name for inp in list(ort_session.get_inputs())]: + ort_inputs[name] = arr + + # output has shape (batch, output_chunk_length, n components, 1 or n likelihood params) + return ort_session.run(None, ort_inputs) diff --git a/requirements/optional.txt b/requirements/optional.txt new file mode 100644 index 0000000000..b36e6c71f8 --- /dev/null +++ b/requirements/optional.txt @@ -0,0 +1,4 @@ +onnx +optuna +ray +onnxruntime From bc70f8193d066741067530c82198bcbf65706f8d Mon Sep 17 00:00:00 2001 From: madtoinou Date: Mon, 24 Feb 2025 19:35:25 +0100 Subject: [PATCH 02/17] fix: typos --- darts/tests/optional_deps/test_onnx.py | 35 +++++++++++++++----------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index 3217a611e2..2d8a988d02 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -14,10 +14,8 @@ ) import onnx import onnxruntime as ort -import torch from darts.models import ( - BlockRNNModel, NLinearModel, TFTModel, TiDEModel, @@ -27,14 +25,15 @@ class TestOnnx: - ts_tgt = tg.linear_timeseries(start_value=0, end_value=100, length=100) - ts_pc = tg.constant_timeseries(value=123.4, length=100) - ts_fc = tg.sine_timeseries(length=100) + ts_tg = tg.linear_timeseries( + start_value=0, end_value=100, length=30, dtype=np.float64 + ) + ts_pc = tg.constant_timeseries(value=123.4, length=30, dtype=np.float64) + ts_fc = tg.sine_timeseries(length=32, dtype=np.float64) @pytest.mark.parametrize( "model_cls", [ - BlockRNNModel, NLinearModel, TFTModel, TiDEModel, @@ -44,8 +43,13 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): model = model_cls( input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs ) - onnx_filename = f"test_{model.name}" - model.fit() + onnx_filename = f"test_{model}" + model.fit( + series=self.ts_tg, + past_covariates=self.ts_pc if model.supports_past_covariates else None, + future_covariates=self.ts_fc if model.supports_future_covariates else None, + ) + # native inference pred = model.predict(2) @@ -59,9 +63,9 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): # manual feature extraction from the series past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( model=model, - series=self.ts_tgt, - past_covariates=self.ts_pc, - future_covariates=self.ts_fc, + series=self.ts_tg, + past_covariates=self.ts_pc if model.supports_past_covariates else None, + future_covariates=self.ts_fc if model.supports_future_covariates else None, ) # onnx model loading and inference @@ -70,9 +74,10 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): ) # check that the predictions are similar - torch.testing.assert_close(onnx_pred, pred) + np.testing.assert_array_almost_equal(onnx_pred[0][0], pred.all_values()) def _helper_prepare_onnx_inputs( + self, model, series: TimeSeries, past_covariates: Optional[TimeSeries] = None, @@ -112,9 +117,9 @@ def _helper_prepare_onnx_inputs( def _helper_onnx_inference( self, onnx_filename: str, - past_feats: torch.Tensor, - future_feats: torch.Tensor, - static_feats: torch.Tensor, + past_feats: np.ndarray, + future_feats: np.ndarray, + static_feats: np.ndarray, ): ort_session = ort.InferenceSession(onnx_filename) # extract only the features expected by the model From 8c72758e8a4c08e295fa354f769eb148cbf52a59 Mon Sep 17 00:00:00 2001 From: madtoinou Date: Tue, 25 Feb 2025 15:58:04 +0100 Subject: [PATCH 03/17] fix: typo --- darts/models/forecasting/tft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/darts/models/forecasting/tft_model.py b/darts/models/forecasting/tft_model.py index 2f53e12af5..a75661be35 100644 --- a/darts/models/forecasting/tft_model.py +++ b/darts/models/forecasting/tft_model.py @@ -726,7 +726,7 @@ def __init__( If ``False``, only attends to previous time steps in the decoder. If ``True`` attends to previous, current, and future time steps. Defaults to ``False``. feed_forward - A feedforward network is a fully-connected layer with an activation. TFT Can be one of the glu variant's + A feedforward network is a fully-connected layer with an activation. Can be one of the glu variant's FeedForward Network (FFN)[2]. The glu variant's FeedForward Network are a series of FFNs designed to work better with Transformer based models. Defaults to ``"GatedResidualNetwork"``. ["GLU", "Bilinear", "ReGLU", "GEGLU", "SwiGLU", "ReLU", "GELU"] or the TFT original FeedForward Network ["GatedResidualNetwork"]. From b1d55d2dfc7cc9c0c9a3b1c94db78246b638db72 Mon Sep 17 00:00:00 2001 From: madtoinou Date: Tue, 25 Feb 2025 15:58:52 +0100 Subject: [PATCH 04/17] fix: tests, extended to more models --- .github/workflows/develop.yml | 4 ++-- .github/workflows/merge.yml | 6 ++--- .github/workflows/update-cache.yml | 4 ++-- darts/tests/optional_deps/test_onnx.py | 33 +++++++++++++++++++------- requirements/optional.txt | 2 +- 5 files changed, 33 insertions(+), 16 deletions(-) diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml index bf62fc285d..7ccf8b8288 100644 --- a/.github/workflows/develop.yml +++ b/.github/workflows/develop.yml @@ -51,7 +51,7 @@ jobs: if [ "${{ matrix.os }}" == "macos-13" ]; then source $HOME/.local/bin/env fi - uv pip compile requirements/dev-all.txt > requirements-latest.txt + uv pip compile requirements/dev-all.txt requirements/optional.txt -o requirements-latest.txt - name: "Cache python environment" uses: actions/cache@v4 @@ -67,7 +67,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt + pip install -U -r requirements/dev-all.txt requirements/optional.txt - name: "Install libomp (for LightGBM)" run: | diff --git a/.github/workflows/merge.yml b/.github/workflows/merge.yml index c4a6e5ce86..6ff978877e 100644 --- a/.github/workflows/merge.yml +++ b/.github/workflows/merge.yml @@ -94,7 +94,7 @@ jobs: - name: "Compile Dependency Versions" run: | curl -LsSf https://astral.sh/uv/install.sh | sh - uv pip compile requirements/dev-all.txt > requirements-latest.txt + uv pip compile requirements/dev-all.txt requirements/optional.txt -o requirements-latest.txt # only restore cache but do not upload - name: "Restore cached python environment" @@ -111,7 +111,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt + pip install -U -r requirements/dev-all.txt requirements/optional.txt - name: "Install libomp (for LightGBM)" run: | @@ -141,7 +141,7 @@ jobs: - name: "Compile Dependency Versions" run: | curl -LsSf https://astral.sh/uv/install.sh | sh - uv pip compile requirements/dev-all.txt > requirements-latest.txt + uv pip compile requirements/dev-all.txt requirements/optional.txt -o requirements-latest.txt # only restore cache but do not upload - name: "Restore cached python environment" diff --git a/.github/workflows/update-cache.yml b/.github/workflows/update-cache.yml index 3243bc4dc2..3b9629a411 100644 --- a/.github/workflows/update-cache.yml +++ b/.github/workflows/update-cache.yml @@ -31,7 +31,7 @@ jobs: if [ "${{ matrix.os }}" == "macos-13" ]; then source $HOME/.local/bin/env fi - uv pip compile requirements/dev-all.txt > requirements-latest.txt + uv pip compile requirements/dev-all.txt requirements/optional.txt -o requirements-latest.txt - name: "Cache python environment" uses: actions/cache@v4 @@ -47,4 +47,4 @@ jobs: - name: "Install Latest Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt + pip install -U -r requirements/dev-all.txt requirements/optional.txt diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index 2d8a988d02..e2fe7bef65 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -16,20 +16,25 @@ import onnxruntime as ort from darts.models import ( + NBEATSModel, + NHiTSModel, NLinearModel, + TCNModel, TFTModel, TiDEModel, + TransformerModel, + TSMixerModel, ) -# from darts.models.components.layer_norm_variants import RINorm +# TODO: check how RINorm can be handled with respect to ONNX class TestOnnx: - ts_tg = tg.linear_timeseries( - start_value=0, end_value=100, length=30, dtype=np.float64 + ts_tg = tg.linear_timeseries(start_value=0, end_value=100, length=30).astype( + "float32" ) - ts_pc = tg.constant_timeseries(value=123.4, length=30, dtype=np.float64) - ts_fc = tg.sine_timeseries(length=32, dtype=np.float64) + ts_pc = tg.constant_timeseries(value=123.4, length=300).astype("float32") + ts_fc = tg.sine_timeseries(length=32).astype("float32") @pytest.mark.parametrize( "model_cls", @@ -37,6 +42,11 @@ class TestOnnx: NLinearModel, TFTModel, TiDEModel, + TCNModel, + NBEATSModel, + NHiTSModel, + TransformerModel, + TSMixerModel, ], ) def test_onnx_save_load(self, tmpdir_fn, model_cls): @@ -44,16 +54,20 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs ) onnx_filename = f"test_{model}" + # model.model = model.model.to(torch.float) model.fit( series=self.ts_tg, past_covariates=self.ts_pc if model.supports_past_covariates else None, future_covariates=self.ts_fc if model.supports_future_covariates else None, ) - + # model.model = model.model.float() # native inference pred = model.predict(2) # model export + # TODO: LSTM model should be exported with a batch size of 1, it seems to create prediction shape problems for + # for TFT and TCN. + model.to_onnx(onnx_filename) # onnx model verification @@ -71,10 +85,11 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): # onnx model loading and inference onnx_pred = self._helper_onnx_inference( onnx_filename, past_feats, future_feats, static_feats - ) + )[0][0] # check that the predictions are similar - np.testing.assert_array_almost_equal(onnx_pred[0][0], pred.all_values()) + assert pred.shape == onnx_pred.shape, "forecasts don't have the same shape." + np.testing.assert_array_almost_equal(onnx_pred, pred.all_values(), decimal=4) def _helper_prepare_onnx_inputs( self, @@ -112,6 +127,8 @@ def _helper_prepare_onnx_inputs( static_feats = np.expand_dims( series.static_covariates_values(), axis=0 ).astype(series.dtype) + + print(series.dtype) return past_feats, future_feats, static_feats def _helper_onnx_inference( diff --git a/requirements/optional.txt b/requirements/optional.txt index b36e6c71f8..b25966fdf0 100644 --- a/requirements/optional.txt +++ b/requirements/optional.txt @@ -1,4 +1,4 @@ onnx +onnxruntime optuna ray -onnxruntime From 3a7a2af7d2a53fd130017a52108e9653f60d476a Mon Sep 17 00:00:00 2001 From: madtoinou Date: Tue, 25 Feb 2025 17:10:45 +0100 Subject: [PATCH 05/17] fix: also cover exporting to onnx after laoding ckpt --- .../forecasting/torch_forecasting_model.py | 4 +- darts/tests/optional_deps/test_onnx.py | 135 +++++++++++++++--- 2 files changed, 115 insertions(+), 24 deletions(-) diff --git a/darts/models/forecasting/torch_forecasting_model.py b/darts/models/forecasting/torch_forecasting_model.py index b949efb914..5dfcc483b2 100644 --- a/darts/models/forecasting/torch_forecasting_model.py +++ b/darts/models/forecasting/torch_forecasting_model.py @@ -677,6 +677,8 @@ def to_onnx(self, path: Optional[str] = None, **kwargs): ``input_sample``, ``input_name``). For more information, read the `official documentation `_. """ + # TODO: LSTM model should be exported with a batch size of 1 + # TODO: predictions with TFT and TCN models is incorrect, might be caused by helper function to process inputs if not self._fit_called: raise_log( ValueError("`fit()` needs to be called before `to_onnx()`."), logger @@ -1774,7 +1776,7 @@ def save( self.trainer.save_checkpoint(path_ptl_ckpt, weights_only=clean) # TODO: keep track of PyTorch Lightning to see if they implement model checkpoint saving - # without having to call fit/predict/validate/test before + # without having to call fit/predict/validate/test before # try to recover original automatic PL checkpoint elif self.load_ckpt_path: if os.path.exists(self.load_ckpt_path): diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index e2fe7bef65..67177836a8 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -1,3 +1,4 @@ +from itertools import product from typing import Optional import numpy as np @@ -16,17 +17,26 @@ import onnxruntime as ort from darts.models import ( + BlockRNNModel, NBEATSModel, NHiTSModel, NLinearModel, - TCNModel, - TFTModel, TiDEModel, TransformerModel, - TSMixerModel, ) # TODO: check how RINorm can be handled with respect to ONNX +torch_model_cls = [ + BlockRNNModel, + NBEATSModel, + NHiTSModel, + NLinearModel, + # TCNModel, + # TFTModel, + TiDEModel, + TransformerModel, + # TSMixerModel, +] class TestOnnx: @@ -36,38 +46,22 @@ class TestOnnx: ts_pc = tg.constant_timeseries(value=123.4, length=300).astype("float32") ts_fc = tg.sine_timeseries(length=32).astype("float32") - @pytest.mark.parametrize( - "model_cls", - [ - NLinearModel, - TFTModel, - TiDEModel, - TCNModel, - NBEATSModel, - NHiTSModel, - TransformerModel, - TSMixerModel, - ], - ) + @pytest.mark.parametrize("model_cls", torch_model_cls) def test_onnx_save_load(self, tmpdir_fn, model_cls): model = model_cls( input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs ) - onnx_filename = f"test_{model}" - # model.model = model.model.to(torch.float) + onnx_filename = f"test_onnx_{model.model_name}.onnx" + model.fit( series=self.ts_tg, past_covariates=self.ts_pc if model.supports_past_covariates else None, future_covariates=self.ts_fc if model.supports_future_covariates else None, ) - # model.model = model.model.float() # native inference pred = model.predict(2) # model export - # TODO: LSTM model should be exported with a batch size of 1, it seems to create prediction shape problems for - # for TFT and TCN. - model.to_onnx(onnx_filename) # onnx model verification @@ -91,6 +85,102 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): assert pred.shape == onnx_pred.shape, "forecasts don't have the same shape." np.testing.assert_array_almost_equal(onnx_pred, pred.all_values(), decimal=4) + @pytest.mark.parametrize( + "params", + product( + torch_model_cls, + [True, False], # clean + ), + ) + def test_onnx_from_ckpt(self, tmpdir_fn, params): + """Check that creating the onnx export from a model directly loaded from a checkpoint work as expected""" + model_cls, clean = params + model = model_cls( + input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs + ) + onnx_filename = f"test_onnx_{model.model_name}.onnx" + onnx_filename2 = f"test_onnx_{model.model_name}_weights.onnx" + ckpt_filename = f"test_ckpt_{model.model_name}.pt" + + model.fit( + series=self.ts_tg, + past_covariates=self.ts_pc if model.supports_past_covariates else None, + future_covariates=self.ts_fc if model.supports_future_covariates else None, + ) + model.save(ckpt_filename, clean=clean) + + # load the entire checkpoint + model_loaded = model_cls.load(ckpt_filename) + pred = model_loaded.predict( + n=2, + series=self.ts_tg, + past_covariates=self.ts_pc + if model_loaded.supports_past_covariates + else None, + future_covariates=self.ts_fc + if model_loaded.supports_future_covariates + else None, + ) + + # export the loaded model + model_loaded.to_onnx(onnx_filename) + + # manual feature extraction from the series + past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( + model=model_loaded, + series=self.ts_tg, + past_covariates=self.ts_pc if model.supports_past_covariates else None, + future_covariates=self.ts_fc if model.supports_future_covariates else None, + ) + + # onnx model loading and inference + onnx_pred = self._helper_onnx_inference( + onnx_filename, past_feats, future_feats, static_feats + )[0][0] + + # check that the predictions are similar + assert pred.shape == onnx_pred.shape, "forecasts don't have the same shape." + np.testing.assert_array_almost_equal(onnx_pred, pred.all_values(), decimal=4) + + # load only the weights + model_weights = model_cls( + input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs + ) + model_weights.load_weights(ckpt_filename) + pred_weights = model_weights.predict( + n=2, + series=self.ts_tg, + past_covariates=self.ts_pc + if model_weights.supports_past_covariates + else None, + future_covariates=self.ts_fc + if model_weights.supports_future_covariates + else None, + ) + + # export the loaded model + model_weights.to_onnx(onnx_filename2) + + # manual feature extraction from the series + past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( + model=model_weights, + series=self.ts_tg, + past_covariates=self.ts_pc if model.supports_past_covariates else None, + future_covariates=self.ts_fc if model.supports_future_covariates else None, + ) + + # onnx model loading and inference + onnx_pred_weights = self._helper_onnx_inference( + onnx_filename2, past_feats, future_feats, static_feats + )[0][0] + + assert pred_weights.shape == onnx_pred_weights.shape, ( + "forecasts don't have the same shape." + ) + np.testing.assert_array_almost_equal( + onnx_pred_weights, pred_weights.all_values(), decimal=4 + ) + def _helper_prepare_onnx_inputs( self, model, @@ -128,7 +218,6 @@ def _helper_prepare_onnx_inputs( series.static_covariates_values(), axis=0 ).astype(series.dtype) - print(series.dtype) return past_feats, future_feats, static_feats def _helper_onnx_inference( From 977e8094fdd90215ddf4e66ada387b1bd8db69bf Mon Sep 17 00:00:00 2001 From: madtoinou Date: Thu, 27 Feb 2025 16:57:41 +0100 Subject: [PATCH 06/17] fix: use_X_covariate attribute is correctly updated after loading weights from checkpoint --- .../forecasting/torch_forecasting_model.py | 57 +++++++++++++++++++ .../test_torch_forecasting_model.py | 57 +++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/darts/models/forecasting/torch_forecasting_model.py b/darts/models/forecasting/torch_forecasting_model.py index 5dfcc483b2..86f7763092 100644 --- a/darts/models/forecasting/torch_forecasting_model.py +++ b/darts/models/forecasting/torch_forecasting_model.py @@ -646,6 +646,12 @@ def _verify_past_future_covariates(self, past_covariates, future_covariates): logger=logger, ) + @abstractmethod + def _update_covariates_use(self): + """Based on the Forecasting class and the training_sample attribute, update the + uses_[past/future/static]_covariates attributes.""" + pass + def to_onnx(self, path: Optional[str] = None, **kwargs): """Export model to ONNX format for optimized inference, wrapping around PyTorch Lightning's :func:`torch.onnx.export` method (`official documentation Date: Thu, 27 Feb 2025 17:07:47 +0100 Subject: [PATCH 07/17] update the test of the onnx optional dep, add a util file --- .../forecasting/torch_forecasting_model.py | 5 -- darts/tests/optional_deps/test_onnx.py | 22 ++++---- darts/utils/onnx_utils.py | 52 +++++++++++++++++++ docs/userguide/torch_forecasting_models.md | 14 ++++- 4 files changed, 74 insertions(+), 19 deletions(-) create mode 100644 darts/utils/onnx_utils.py diff --git a/darts/models/forecasting/torch_forecasting_model.py b/darts/models/forecasting/torch_forecasting_model.py index 86f7763092..3a003bcf25 100644 --- a/darts/models/forecasting/torch_forecasting_model.py +++ b/darts/models/forecasting/torch_forecasting_model.py @@ -2696,7 +2696,6 @@ def extreme_lags( def _update_covariates_use(self): """The model is expected to rely on the `PastCovariatesTrainingDataset`""" - print(self.train_sample) _, past_covs, static_covs, _ = self.train_sample self._uses_past_covariates = past_covs is not None self._uses_future_covariates = False @@ -2797,7 +2796,6 @@ def extreme_lags( def _update_covariates_use(self): """The model is expected to rely on the `FutureCovariatesTrainingDataset`""" - print(self.train_sample) _, future_covs, static_covs, _ = self.train_sample self._uses_past_covariates = False self._uses_future_covariates = future_covs is not None @@ -2899,7 +2897,6 @@ def extreme_lags( def _update_covariates_use(self): """The model is expected to rely on the `DualCovariatesTrainingDataset`""" - print(self.train_sample) _, historic_future_covs, future_covs, static_covs, _ = self.train_sample self._uses_past_covariates = False self._uses_future_covariates = ( @@ -3003,7 +3000,6 @@ def extreme_lags( def _update_covariates_use(self): """The model is expected to rely on the `MixedCovariatesTrainingDataset`""" - print(self.train_sample) _, past_covs, historic_future_covs, future_covs, static_covs, _ = ( self.train_sample ) @@ -3110,7 +3106,6 @@ def extreme_lags( def _update_covariates_use(self): """The model is expected to rely on the `SplitCovariatesTrainingDataset`""" - print(self.train_sample) _, past_covs, historic_future_covs, future_covs, static_covs, _ = ( self.train_sample ) diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index 67177836a8..6ddfec7c0c 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -72,8 +72,8 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( model=model, series=self.ts_tg, - past_covariates=self.ts_pc if model.supports_past_covariates else None, - future_covariates=self.ts_fc if model.supports_future_covariates else None, + past_covariates=self.ts_pc if model.uses_past_covariates else None, + future_covariates=self.ts_fc if model.uses_future_covariates else None, ) # onnx model loading and inference @@ -114,11 +114,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): pred = model_loaded.predict( n=2, series=self.ts_tg, - past_covariates=self.ts_pc - if model_loaded.supports_past_covariates - else None, + past_covariates=self.ts_pc if model_loaded.uses_past_covariates else None, future_covariates=self.ts_fc - if model_loaded.supports_future_covariates + if model_loaded.uses_future_covariates else None, ) @@ -129,8 +127,10 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( model=model_loaded, series=self.ts_tg, - past_covariates=self.ts_pc if model.supports_past_covariates else None, - future_covariates=self.ts_fc if model.supports_future_covariates else None, + past_covariates=self.ts_pc if model_loaded.uses_past_covariates else None, + future_covariates=self.ts_fc + if model_loaded.uses_future_covariates + else None, ) # onnx model loading and inference @@ -150,11 +150,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): pred_weights = model_weights.predict( n=2, series=self.ts_tg, - past_covariates=self.ts_pc - if model_weights.supports_past_covariates - else None, + past_covariates=self.ts_pc if model_weights.uses_past_covariates else None, future_covariates=self.ts_fc - if model_weights.supports_future_covariates + if model_weights.uses_future_covariates else None, ) diff --git a/darts/utils/onnx_utils.py b/darts/utils/onnx_utils.py new file mode 100644 index 0000000000..39d3b67c33 --- /dev/null +++ b/darts/utils/onnx_utils.py @@ -0,0 +1,52 @@ +from typing import Optional + +import numpy as np + +from darts import TimeSeries + + +def prepare_onnx_inputs( + model, + series: TimeSeries, + past_covariates: Optional[TimeSeries] = None, + future_covariates: Optional[TimeSeries] = None, +) -> tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]: + """Helper function to slice and concatenate the input features. + + In order to remove the dependency on the `model` argument, it can be decomposed into + the following arguments (and simplified depending on the characteristics of the model used): + - model_icl + - model_ocl + - model_uses_past_covs + - model_uses_future_covs + - model_uses_static_covs + """ + past_feats, future_feats, static_feats = None, None, None + # get input & output windows + past_start = series.end_time() - (model.input_chunk_length - 1) * series.freq + past_end = series.end_time() + future_start = past_end + 1 * series.freq + future_end = past_end + model.output_chunk_length * series.freq + # extract all historic and future features from target, past and future covariates + past_feats = series[past_start:past_end].values() + if past_covariates and model.uses_past_covariates: + # extract past covariates + past_feats = np.concatenate( + [past_feats, past_covariates[past_start:past_end].values()], axis=1 + ) + if future_covariates and model.uses_future_covariates: + # extract past part of future covariates + past_feats = np.concatenate( + [past_feats, future_covariates[past_start:past_end].values()], axis=1 + ) + # extract future part of future covariates + future_feats = future_covariates[future_start:future_end].values() + # add batch dimension -> (batch, n time steps, n components) + past_feats = np.expand_dims(past_feats, axis=0).astype(series.dtype) + future_feats = np.expand_dims(future_feats, axis=0).astype(series.dtype) + # extract static covariates + if series.has_static_covariates and model.uses_static_covariates: + static_feats = np.expand_dims(series.static_covariates_values(), axis=0).astype( + series.dtype + ) + return past_feats, future_feats, static_feats diff --git a/docs/userguide/torch_forecasting_models.md b/docs/userguide/torch_forecasting_models.md index 928612e7f5..8b65345bd0 100644 --- a/docs/userguide/torch_forecasting_models.md +++ b/docs/userguide/torch_forecasting_models.md @@ -373,13 +373,23 @@ import onnxruntime as ort import numpy as np from darts import TimeSeries +# can be imported from darts.utils.onnx_utils.py def prepare_onnx_inputs( model, series: TimeSeries, past_covariates : Optional[TimeSeries] = None, future_covariates : Optional[TimeSeries] = None, -) -> tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: - """Helper function to slice and concatenate the input features""" +) -> tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]: + """Helper function to slice and concatenate the input features. + + In order to remove the dependency on the `model` argument, it can be decomposed into + the following arguments (and simplified depending on the characteristics of the model used): + - model_icl + - model_ocl + - model_uses_past_covs + - model_uses_future_covs + - model_uses_static_covs + """ past_feats, future_feats, static_feats = None, None, None # get input & output windows past_start = series.end_time() - (model.input_chunk_length - 1) * series.freq From 7400db8a09cfc14983c33004b8d4e492162573b4 Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 10:19:02 +0100 Subject: [PATCH 08/17] reduced code redundancy in the tests --- darts/tests/optional_deps/test_onnx.py | 87 +++++++------------------- 1 file changed, 24 insertions(+), 63 deletions(-) diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index 6ddfec7c0c..b42b71d1b1 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -7,6 +7,7 @@ import darts.utils.timeseries_generation as tg from darts import TimeSeries from darts.tests.conftest import ONNX_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs +from darts.utils.onnx_utils import prepare_onnx_inputs if not (TORCH_AVAILABLE and ONNX_AVAILABLE): pytest.skip( @@ -68,17 +69,13 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): onnx_model = onnx.load(onnx_filename) onnx.checker.check_model(onnx_model) - # manual feature extraction from the series - past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( + # onnx model loading and inference + onnx_pred = self._helper_onnx_inference( model=model, + onnx_filename=onnx_filename, series=self.ts_tg, past_covariates=self.ts_pc if model.uses_past_covariates else None, future_covariates=self.ts_fc if model.uses_future_covariates else None, - ) - - # onnx model loading and inference - onnx_pred = self._helper_onnx_inference( - onnx_filename, past_feats, future_feats, static_feats )[0][0] # check that the predictions are similar @@ -123,19 +120,15 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): # export the loaded model model_loaded.to_onnx(onnx_filename) - # manual feature extraction from the series - past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( + # onnx model loading and inference + onnx_pred = self._helper_onnx_inference( model=model_loaded, + onnx_filename=onnx_filename, series=self.ts_tg, past_covariates=self.ts_pc if model_loaded.uses_past_covariates else None, future_covariates=self.ts_fc if model_loaded.uses_future_covariates else None, - ) - - # onnx model loading and inference - onnx_pred = self._helper_onnx_inference( - onnx_filename, past_feats, future_feats, static_feats )[0][0] # check that the predictions are similar @@ -159,17 +152,13 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): # export the loaded model model_weights.to_onnx(onnx_filename2) - # manual feature extraction from the series - past_feats, future_feats, static_feats = self._helper_prepare_onnx_inputs( + # onnx model loading and inference + onnx_pred_weights = self._helper_onnx_inference( model=model_weights, + onnx_filename=onnx_filename2, series=self.ts_tg, past_covariates=self.ts_pc if model.supports_past_covariates else None, future_covariates=self.ts_fc if model.supports_future_covariates else None, - ) - - # onnx model loading and inference - onnx_pred_weights = self._helper_onnx_inference( - onnx_filename2, past_feats, future_feats, static_feats )[0][0] assert pred_weights.shape == onnx_pred_weights.shape, ( @@ -179,53 +168,25 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): onnx_pred_weights, pred_weights.all_values(), decimal=4 ) - def _helper_prepare_onnx_inputs( - self, - model, - series: TimeSeries, - past_covariates: Optional[TimeSeries] = None, - future_covariates: Optional[TimeSeries] = None, - ) -> tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: - """Helper function to slice and concatenate the input features""" - past_feats, future_feats, static_feats = None, None, None - # get input & output windows - past_start = series.end_time() - (model.input_chunk_length - 1) * series.freq - past_end = series.end_time() - future_start = past_end + 1 * series.freq - future_end = past_end + model.output_chunk_length * series.freq - # extract all historic and future features from target, past and future covariates - past_feats = series[past_start:past_end].values() - if past_covariates and model.uses_past_covariates: - # extract past covariates - past_feats = np.concatenate( - [past_feats, past_covariates[past_start:past_end].values()], axis=1 - ) - if future_covariates and model.uses_future_covariates: - # extract past part of future covariates - past_feats = np.concatenate( - [past_feats, future_covariates[past_start:past_end].values()], axis=1 - ) - # extract future part of future covariates - future_feats = future_covariates[future_start:future_end].values() - # add batch dimension -> (batch, n time steps, n components) - past_feats = np.expand_dims(past_feats, axis=0).astype(series.dtype) - future_feats = np.expand_dims(future_feats, axis=0).astype(series.dtype) - # extract static covariates - if series.has_static_covariates and model.uses_static_covariates: - static_feats = np.expand_dims( - series.static_covariates_values(), axis=0 - ).astype(series.dtype) - - return past_feats, future_feats, static_feats - def _helper_onnx_inference( self, + model, onnx_filename: str, - past_feats: np.ndarray, - future_feats: np.ndarray, - static_feats: np.ndarray, + series: TimeSeries, + past_covariates: Optional[TimeSeries], + future_covariates: Optional[TimeSeries], ): + """Darts model is only used to detect which covariates are supported by the weights.""" ort_session = ort.InferenceSession(onnx_filename) + + # extract the input arrays from the series + past_feats, future_feats, static_feats = prepare_onnx_inputs( + model=model, + series=series, + past_covariates=past_covariates, + future_covariates=future_covariates, + ) + # extract only the features expected by the model ort_inputs = {} for name, arr in zip( From f783e1a601878109d984010b1c19082df55f21b6 Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 11:14:36 +0100 Subject: [PATCH 09/17] feat: adding tests for optuna --- darts/tests/conftest.py | 16 ++ darts/tests/optional_deps/test_optuna.py | 188 +++++++++++++++++++++++ requirements/optional.txt | 1 + 3 files changed, 205 insertions(+) create mode 100644 darts/tests/optional_deps/test_optuna.py diff --git a/darts/tests/conftest.py b/darts/tests/conftest.py index 29059d0bf9..3c7288b525 100644 --- a/darts/tests/conftest.py +++ b/darts/tests/conftest.py @@ -26,6 +26,22 @@ logger.warning("Onnx not installed - Some tests will be skipped.") ONNX_AVAILABLE = False +try: + import optuna # noqa: F401 + + OPTUNA_AVAILABLE = True +except ImportError: + logger.warning("Optuna not installed - Some tests will be skipped.") + OPTUNA_AVAILABLE = False + +try: + import ray # noqa: F401 + + RAY_AVAILABLE = True +except ImportError: + logger.warning("Ray not installed - Some tests will be skipped.") + RAY_AVAILABLE = False + tfm_kwargs = { "pl_trainer_kwargs": { "accelerator": "cpu", diff --git a/darts/tests/optional_deps/test_optuna.py b/darts/tests/optional_deps/test_optuna.py new file mode 100644 index 0000000000..dfed1dd21c --- /dev/null +++ b/darts/tests/optional_deps/test_optuna.py @@ -0,0 +1,188 @@ +import os +from itertools import product + +import numpy as np +import pytest +from sklearn.preprocessing import MaxAbsScaler + +from darts.dataprocessing.transformers import Scaler +from darts.datasets import AirPassengersDataset +from darts.metrics import smape +from darts.models import LinearRegressionModel +from darts.tests.conftest import OPTUNA_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs + +if not OPTUNA_AVAILABLE: + pytest.skip( + f"Optuna not available. {__name__} tests will be skipped.", + allow_module_level=True, + ) + +import optuna + +if TORCH_AVAILABLE: + import torch + from pytorch_lightning.callbacks import Callback, EarlyStopping + + # hacky workaround found in https://github.com/Lightning-AI/pytorch-lightning/issues/17485 + # to avoid import of both lightning and pytorch_lightning + class PatchedCallback(optuna.integration.PyTorchLightningPruningCallback, Callback): + pass + + from darts.models import TCNModel + from darts.utils.likelihood_models import GaussianLikelihood + + +class TestOptuna: + series = AirPassengersDataset().load().astype(np.float32) + + val_length = 36 + train, val = series.split_after(val_length) + + # scale + scaler = Scaler(MaxAbsScaler()) + train = scaler.fit_transform(train) + val = scaler.transform(val) + + @pytest.mark.skipif(not TORCH_AVAILABLE, reason="requires torch") + def test_hp_opti_torch_model(self, tmpdir_fn): + """Check that optuna works as expected with a torch-based model""" + + # define objective function + def objective(trial): + # select input and output chunk lengths + in_len = trial.suggest_int("in_len", 4, 8) + out_len = trial.suggest_int("out_len", 1, 3) + + # Other hyperparameters + kernel_size = trial.suggest_int("kernel_size", 2, 3) + num_filters = trial.suggest_int("num_filters", 1, 2) + lr = trial.suggest_float("lr", 5e-5, 1e-3, log=True) + include_year = trial.suggest_categorical("year", [False, True]) + + # throughout training we'll monitor the validation loss for both pruning and early stopping + pruner = PatchedCallback(trial, monitor="val_loss") + early_stopper = EarlyStopping( + "val_loss", min_delta=0.001, patience=3, verbose=True + ) + + pl_trainer_kwargs = { + **tfm_kwargs["pl_trainer_kwargs"], + "callbacks": [pruner, early_stopper], + } + + # optionally also add the (scaled) year value as a past covariate + if include_year: + encoders = { + "datetime_attribute": {"past": ["year"]}, + "transformer": Scaler(), + } + else: + encoders = None + + # reproducibility + torch.manual_seed(42) + + # build the TCN model + model = TCNModel( + input_chunk_length=in_len, + output_chunk_length=out_len, + batch_size=8, + n_epochs=1, + nr_epochs_val_period=1, + kernel_size=kernel_size, + num_filters=num_filters, + optimizer_kwargs={"lr": lr}, + add_encoders=encoders, + likelihood=GaussianLikelihood(), + pl_trainer_kwargs=pl_trainer_kwargs, + model_name="tcn_model", + force_reset=True, + save_checkpoints=True, + work_dir=os.getcwd(), + ) + + # when validating during training, we can use a slightly longer validation + # set which also contains the first input_chunk_length time steps + model_val_set = self.scaler.transform( + self.series[-(self.val_length + in_len) :] + ) + + # train the model + model.fit( + series=self.train, + val_series=model_val_set, + ) + + # reload best model over course of training + model = TCNModel.load_from_checkpoint( + model_name="tcn_model", work_dir=os.getcwd() + ) + + # Evaluate how good it is on the validation set, using sMAPE + preds = model.predict(series=self.train, n=self.val_length) + smapes = smape(self.val, preds, n_jobs=-1) + smape_val = np.mean(smapes) + + return smape_val if smape_val != np.nan else float("inf") + + # optimize hyperparameters by minimizing the sMAPE on the validation set + study = optuna.create_study(direction="minimize") + study.optimize(objective, n_trials=4) + + @pytest.mark.parametrize( + "params", + product( + [True, False], # multi_models + [1, 3], # ocl + ), + ) + def test_hp_opti_regression_model(self, params): + """Check that optuna works as expected with a regression model""" + + multi_models, ocl = params + + # define objective function + def objective(trial): + # select input and output chunk lengths + target_lags = trial.suggest_int("lags", 1, 12) + + include_year = trial.suggest_categorical("year", [False, True]) + + # optionally also add the (scaled) year value as a past covariate + if include_year: + encoders = { + "datetime_attribute": {"past": ["year"]}, + "transformer": Scaler(), + } + past_lags = trial.suggest_int("lags_past_covariates", 1, 12) + else: + encoders = None + past_lags = None + + # reproducibility + torch.manual_seed(42) + + # build the TCN model + model = LinearRegressionModel( + lags=target_lags, + lags_past_covariates=past_lags, + output_chunk_length=ocl, + multi_models=multi_models, + add_encoders=encoders, + ) + + # train the model + model.fit( + series=self.train, + ) + + # Evaluate how good it is on the validation set, using sMAPE + preds = model.predict(series=self.train, n=self.val_length) + smapes = smape(self.val, preds, n_jobs=-1) + smape_val = np.mean(smapes) + + return smape_val if smape_val != np.nan else float("inf") + + # optimize hyperparameters by minimizing the sMAPE on the validation set + study = optuna.create_study(direction="minimize") + study.optimize(objective, n_trials=4) diff --git a/requirements/optional.txt b/requirements/optional.txt index b25966fdf0..b5ba05815a 100644 --- a/requirements/optional.txt +++ b/requirements/optional.txt @@ -1,4 +1,5 @@ onnx onnxruntime optuna +optuna-integration[pytorch_lightning] ray From 83e296295f7700b234a39fed59af4f7b95332beb Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 11:50:07 +0100 Subject: [PATCH 10/17] feat: adding tests for ray --- darts/tests/optional_deps/test_optuna.py | 30 ++-- darts/tests/optional_deps/test_ray.py | 168 +++++++++++++++++++++++ 2 files changed, 179 insertions(+), 19 deletions(-) create mode 100644 darts/tests/optional_deps/test_ray.py diff --git a/darts/tests/optional_deps/test_optuna.py b/darts/tests/optional_deps/test_optuna.py index dfed1dd21c..41368ba793 100644 --- a/darts/tests/optional_deps/test_optuna.py +++ b/darts/tests/optional_deps/test_optuna.py @@ -44,7 +44,7 @@ class TestOptuna: val = scaler.transform(val) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="requires torch") - def test_hp_opti_torch_model(self, tmpdir_fn): + def test_optuna_torch_model(self, tmpdir_fn): """Check that optuna works as expected with a torch-based model""" # define objective function @@ -65,11 +65,6 @@ def objective(trial): "val_loss", min_delta=0.001, patience=3, verbose=True ) - pl_trainer_kwargs = { - **tfm_kwargs["pl_trainer_kwargs"], - "callbacks": [pruner, early_stopper], - } - # optionally also add the (scaled) year value as a past covariate if include_year: encoders = { @@ -87,14 +82,17 @@ def objective(trial): input_chunk_length=in_len, output_chunk_length=out_len, batch_size=8, - n_epochs=1, + n_epochs=2, nr_epochs_val_period=1, kernel_size=kernel_size, num_filters=num_filters, optimizer_kwargs={"lr": lr}, add_encoders=encoders, likelihood=GaussianLikelihood(), - pl_trainer_kwargs=pl_trainer_kwargs, + pl_trainer_kwargs={ + **tfm_kwargs["pl_trainer_kwargs"], + "callbacks": [pruner, early_stopper], + }, model_name="tcn_model", force_reset=True, save_checkpoints=True, @@ -127,7 +125,7 @@ def objective(trial): # optimize hyperparameters by minimizing the sMAPE on the validation set study = optuna.create_study(direction="minimize") - study.optimize(objective, n_trials=4) + study.optimize(objective, n_trials=3) @pytest.mark.parametrize( "params", @@ -136,16 +134,15 @@ def objective(trial): [1, 3], # ocl ), ) - def test_hp_opti_regression_model(self, params): + def test_optuna_regression_model(self, params): """Check that optuna works as expected with a regression model""" multi_models, ocl = params # define objective function def objective(trial): - # select input and output chunk lengths + # select input and encoder usage target_lags = trial.suggest_int("lags", 1, 12) - include_year = trial.suggest_categorical("year", [False, True]) # optionally also add the (scaled) year value as a past covariate @@ -159,10 +156,7 @@ def objective(trial): encoders = None past_lags = None - # reproducibility - torch.manual_seed(42) - - # build the TCN model + # build the model model = LinearRegressionModel( lags=target_lags, lags_past_covariates=past_lags, @@ -170,8 +164,6 @@ def objective(trial): multi_models=multi_models, add_encoders=encoders, ) - - # train the model model.fit( series=self.train, ) @@ -185,4 +177,4 @@ def objective(trial): # optimize hyperparameters by minimizing the sMAPE on the validation set study = optuna.create_study(direction="minimize") - study.optimize(objective, n_trials=4) + study.optimize(objective, n_trials=3) diff --git a/darts/tests/optional_deps/test_ray.py b/darts/tests/optional_deps/test_ray.py new file mode 100644 index 0000000000..7e59a66588 --- /dev/null +++ b/darts/tests/optional_deps/test_ray.py @@ -0,0 +1,168 @@ +import numpy as np +import pytest +from sklearn.preprocessing import MaxAbsScaler + +from darts.dataprocessing.transformers import Scaler +from darts.datasets import AirPassengersDataset +from darts.metrics import smape +from darts.models import LinearRegressionModel +from darts.tests.conftest import RAY_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs + +if not RAY_AVAILABLE: + pytest.skip( + f"Ray not available. {__name__} tests will be skipped.", + allow_module_level=True, + ) + +from ray import tune +from ray.tune.tuner import Tuner + +if TORCH_AVAILABLE: + from pytorch_lightning.callbacks import Callback, EarlyStopping + from ray.tune.integration.pytorch_lightning import TuneReportCheckpointCallback + from torchmetrics import ( + MeanAbsoluteError, + MeanAbsolutePercentageError, + MetricCollection, + ) + + from darts.models import NBEATSModel + + +class TestRay: + series = AirPassengersDataset().load().astype(np.float32) + + val_length = 36 + train, val = series.split_after(val_length) + + # scale + scaler = Scaler(MaxAbsScaler()) + train = scaler.fit_transform(train) + val = scaler.transform(val) + + @pytest.mark.skipif(not TORCH_AVAILABLE, reason="requires torch") + def test_ray_torch_model(self, tmpdir_fn): + """Check that ray works as expected with a torch-based model""" + + def train_model(model_args, callbacks, train, val): + torch_metrics = MetricCollection([ + MeanAbsolutePercentageError(), + MeanAbsoluteError(), + ]) + + # Create the model using model_args from Ray Tune + model = NBEATSModel( + input_chunk_length=4, + output_chunk_length=3, + n_epochs=2, + torch_metrics=torch_metrics, + pl_trainer_kwargs={ + **tfm_kwargs["pl_trainer_kwargs"], + "callbacks": callbacks, + }, + **model_args, + ) + + model.fit( + series=train, + val_series=val, + ) + + # Early stop callback + my_stopper = EarlyStopping( + monitor="val_MeanAbsolutePercentageError", + patience=5, + min_delta=0.05, + mode="min", + ) + + # set up ray tune callback + class TuneReportCallback(TuneReportCheckpointCallback, Callback): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + tune_callback = TuneReportCallback( + { + "loss": "val_loss", + "MAPE": "val_MeanAbsolutePercentageError", + }, + on="validation_end", + ) + + # Define the trainable function that will be tuned by Ray Tune + train_fn_with_parameters = tune.with_parameters( + train_model, + callbacks=[tune_callback, my_stopper], + train=self.train, + val=self.val, + ) + + # Set the resources to be used for each trial (disable GPU, if you don't have one) + resources_per_trial = {"cpu": 2, "gpu": 0} + + # define the hyperparameter space + config = { + "batch_size": tune.choice([8, 16]), + "num_blocks": tune.choice([1, 2]), + "num_stacks": tune.choice([2, 4]), + } + + # the number of combinations to try + num_samples = 2 + + # Create the Tuner object and run the hyperparameter search + tuner = Tuner( + trainable=tune.with_resources( + train_fn_with_parameters, resources=resources_per_trial + ), + param_space=config, + tune_config=tune.TuneConfig( + metric="MAPE", mode="min", num_samples=num_samples + ), + run_config=tune.RunConfig(name="tune_darts"), + ) + tuner.fit() + return + + def test_ray_regression_model(self, tmpdir_fn): + """Check that ray works as expected with a regression model""" + + # define objective function + def objective(config): + # optionally also add the (scaled) year value as a past covariate + if config["include_year"]: + encoders = { + "datetime_attribute": {"past": ["year"]}, + "transformer": Scaler(), + } + past_lags = 1 + else: + encoders = None + past_lags = None + + # build the model + model = LinearRegressionModel( + lags=config["lags"], + lags_past_covariates=past_lags, + output_chunk_length=1, + add_encoders=encoders, + ) + model.fit( + series=self.train, + ) + + # Evaluate how good it is on the validation set, using sMAPE + preds = model.predict(series=self.train, n=self.val_length) + smapes = smape(self.val, preds, n_jobs=-1) + smape_val = np.mean(smapes) + + return smape_val if smape_val != np.nan else float("inf") + + search_space = { + "lags": tune.choice([1, 3, 6, 12]), + "include_years": tune.choice([True, False]), + } + + tuner = tune.Tuner(objective, param_space=search_space) + tuner.fit() + return From d10094deaa7b15a973baeccce71b8171fb294e04 Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 11:54:26 +0100 Subject: [PATCH 11/17] fix: simplified test --- darts/tests/optional_deps/test_ray.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/darts/tests/optional_deps/test_ray.py b/darts/tests/optional_deps/test_ray.py index 7e59a66588..42f867ea28 100644 --- a/darts/tests/optional_deps/test_ray.py +++ b/darts/tests/optional_deps/test_ray.py @@ -78,8 +78,7 @@ def train_model(model_args, callbacks, train, val): # set up ray tune callback class TuneReportCallback(TuneReportCheckpointCallback, Callback): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + pass tune_callback = TuneReportCallback( { @@ -97,11 +96,8 @@ def __init__(self, *args, **kwargs): val=self.val, ) - # Set the resources to be used for each trial (disable GPU, if you don't have one) - resources_per_trial = {"cpu": 2, "gpu": 0} - # define the hyperparameter space - config = { + param_space = { "batch_size": tune.choice([8, 16]), "num_blocks": tune.choice([1, 2]), "num_stacks": tune.choice([2, 4]), @@ -112,17 +108,14 @@ def __init__(self, *args, **kwargs): # Create the Tuner object and run the hyperparameter search tuner = Tuner( - trainable=tune.with_resources( - train_fn_with_parameters, resources=resources_per_trial - ), - param_space=config, + trainable=train_fn_with_parameters, + param_space=param_space, tune_config=tune.TuneConfig( metric="MAPE", mode="min", num_samples=num_samples ), run_config=tune.RunConfig(name="tune_darts"), ) tuner.fit() - return def test_ray_regression_model(self, tmpdir_fn): """Check that ray works as expected with a regression model""" @@ -165,4 +158,3 @@ def objective(config): tuner = tune.Tuner(objective, param_space=search_space) tuner.fit() - return From b9c4e3d92df5c592f9ed1fe635542638fa3e1580 Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 12:07:29 +0100 Subject: [PATCH 12/17] fix: github actions --- .github/workflows/develop.yml | 2 +- .github/workflows/merge.yml | 4 ++-- .github/workflows/update-cache.yml | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml index 7ccf8b8288..27773cfcd1 100644 --- a/.github/workflows/develop.yml +++ b/.github/workflows/develop.yml @@ -67,7 +67,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt requirements/optional.txt + pip install -U -r requirements/dev-all.txt -r requirements/optional.txt - name: "Install libomp (for LightGBM)" run: | diff --git a/.github/workflows/merge.yml b/.github/workflows/merge.yml index 6ff978877e..22bccb81f4 100644 --- a/.github/workflows/merge.yml +++ b/.github/workflows/merge.yml @@ -111,7 +111,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt requirements/optional.txt + pip install -U -r requirements/dev-all.txt - name: "Install libomp (for LightGBM)" run: | @@ -141,7 +141,7 @@ jobs: - name: "Compile Dependency Versions" run: | curl -LsSf https://astral.sh/uv/install.sh | sh - uv pip compile requirements/dev-all.txt requirements/optional.txt -o requirements-latest.txt + uv pip compile requirements/dev-all.txt -o requirements-latest.txt # only restore cache but do not upload - name: "Restore cached python environment" diff --git a/.github/workflows/update-cache.yml b/.github/workflows/update-cache.yml index 3b9629a411..85d03c6f8c 100644 --- a/.github/workflows/update-cache.yml +++ b/.github/workflows/update-cache.yml @@ -31,7 +31,7 @@ jobs: if [ "${{ matrix.os }}" == "macos-13" ]; then source $HOME/.local/bin/env fi - uv pip compile requirements/dev-all.txt requirements/optional.txt -o requirements-latest.txt + uv pip compile requirements/dev-all.txt -r requirements/optional.txt -o requirements-latest.txt - name: "Cache python environment" uses: actions/cache@v4 @@ -47,4 +47,4 @@ jobs: - name: "Install Latest Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt requirements/optional.txt + pip install -U -r requirements/dev-all.txt -r requirements/optional.txt From cbbe27e9affe80e1270c7f865518b537103b145c Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 13:01:57 +0100 Subject: [PATCH 13/17] tmp fix: remove ray test for regression model --- darts/tests/optional_deps/test_ray.py | 44 --------------------------- 1 file changed, 44 deletions(-) diff --git a/darts/tests/optional_deps/test_ray.py b/darts/tests/optional_deps/test_ray.py index 42f867ea28..daf2eaa056 100644 --- a/darts/tests/optional_deps/test_ray.py +++ b/darts/tests/optional_deps/test_ray.py @@ -4,8 +4,6 @@ from darts.dataprocessing.transformers import Scaler from darts.datasets import AirPassengersDataset -from darts.metrics import smape -from darts.models import LinearRegressionModel from darts.tests.conftest import RAY_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs if not RAY_AVAILABLE: @@ -116,45 +114,3 @@ class TuneReportCallback(TuneReportCheckpointCallback, Callback): run_config=tune.RunConfig(name="tune_darts"), ) tuner.fit() - - def test_ray_regression_model(self, tmpdir_fn): - """Check that ray works as expected with a regression model""" - - # define objective function - def objective(config): - # optionally also add the (scaled) year value as a past covariate - if config["include_year"]: - encoders = { - "datetime_attribute": {"past": ["year"]}, - "transformer": Scaler(), - } - past_lags = 1 - else: - encoders = None - past_lags = None - - # build the model - model = LinearRegressionModel( - lags=config["lags"], - lags_past_covariates=past_lags, - output_chunk_length=1, - add_encoders=encoders, - ) - model.fit( - series=self.train, - ) - - # Evaluate how good it is on the validation set, using sMAPE - preds = model.predict(series=self.train, n=self.val_length) - smapes = smape(self.val, preds, n_jobs=-1) - smape_val = np.mean(smapes) - - return smape_val if smape_val != np.nan else float("inf") - - search_space = { - "lags": tune.choice([1, 3, 6, 12]), - "include_years": tune.choice([True, False]), - } - - tuner = tune.Tuner(objective, param_space=search_space) - tuner.fit() From c51516884616db06a7d02b2fe5c6a9348e0f39c6 Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 13:46:06 +0100 Subject: [PATCH 14/17] fix: improve test coverage --- darts/tests/optional_deps/test_onnx.py | 42 ++++++++++++++++++++------ 1 file changed, 33 insertions(+), 9 deletions(-) diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index b42b71d1b1..8d769db157 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -2,6 +2,7 @@ from typing import Optional import numpy as np +import pandas as pd import pytest import darts.utils.timeseries_generation as tg @@ -44,6 +45,9 @@ class TestOnnx: ts_tg = tg.linear_timeseries(start_value=0, end_value=100, length=30).astype( "float32" ) + ts_tg_with_static = ts_tg.with_static_covariates( + pd.Series(data=[12], index=["loc"]) + ) ts_pc = tg.constant_timeseries(value=123.4, length=300).astype("float32") ts_fc = tg.sine_timeseries(length=32).astype("float32") @@ -54,8 +58,14 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): ) onnx_filename = f"test_onnx_{model.model_name}.onnx" + # exporting without fitting the model fails + with pytest.raises(ValueError): + model.to_onnx("dummy_name.onnx") + model.fit( - series=self.ts_tg, + series=self.ts_tg_with_static + if model.supports_static_covariates + else self.ts_tg, past_covariates=self.ts_pc if model.supports_past_covariates else None, future_covariates=self.ts_fc if model.supports_future_covariates else None, ) @@ -73,7 +83,9 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): onnx_pred = self._helper_onnx_inference( model=model, onnx_filename=onnx_filename, - series=self.ts_tg, + series=self.ts_tg_with_static + if model.uses_static_covariates + else self.ts_tg, past_covariates=self.ts_pc if model.uses_past_covariates else None, future_covariates=self.ts_fc if model.uses_future_covariates else None, )[0][0] @@ -100,7 +112,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): ckpt_filename = f"test_ckpt_{model.model_name}.pt" model.fit( - series=self.ts_tg, + series=self.ts_tg_with_static + if model.supports_static_covariates + else self.ts_tg, past_covariates=self.ts_pc if model.supports_past_covariates else None, future_covariates=self.ts_fc if model.supports_future_covariates else None, ) @@ -110,7 +124,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): model_loaded = model_cls.load(ckpt_filename) pred = model_loaded.predict( n=2, - series=self.ts_tg, + series=self.ts_tg_with_static + if model_loaded.uses_static_covariates + else self.ts_tg, past_covariates=self.ts_pc if model_loaded.uses_past_covariates else None, future_covariates=self.ts_fc if model_loaded.uses_future_covariates @@ -124,7 +140,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): onnx_pred = self._helper_onnx_inference( model=model_loaded, onnx_filename=onnx_filename, - series=self.ts_tg, + series=self.ts_tg_with_static + if model_loaded.uses_static_covariates + else self.ts_tg, past_covariates=self.ts_pc if model_loaded.uses_past_covariates else None, future_covariates=self.ts_fc if model_loaded.uses_future_covariates @@ -142,7 +160,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): model_weights.load_weights(ckpt_filename) pred_weights = model_weights.predict( n=2, - series=self.ts_tg, + series=self.ts_tg_with_static + if model_weights.uses_static_covariates + else self.ts_tg, past_covariates=self.ts_pc if model_weights.uses_past_covariates else None, future_covariates=self.ts_fc if model_weights.uses_future_covariates @@ -156,9 +176,13 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): onnx_pred_weights = self._helper_onnx_inference( model=model_weights, onnx_filename=onnx_filename2, - series=self.ts_tg, - past_covariates=self.ts_pc if model.supports_past_covariates else None, - future_covariates=self.ts_fc if model.supports_future_covariates else None, + series=self.ts_tg_with_static + if model_weights.uses_static_covariates + else self.ts_tg, + past_covariates=self.ts_pc if model_weights.uses_past_covariates else None, + future_covariates=self.ts_fc + if model_weights.uses_future_covariates + else None, )[0][0] assert pred_weights.shape == onnx_pred_weights.shape, ( From 624dfdebd4ab5fc94c9de71710bdabf43f17c56d Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 28 Feb 2025 13:49:04 +0100 Subject: [PATCH 15/17] fix: further simply the tests --- darts/tests/optional_deps/test_onnx.py | 28 +++++++++----------------- 1 file changed, 9 insertions(+), 19 deletions(-) diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index 8d769db157..55b7d02091 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -83,11 +83,9 @@ def test_onnx_save_load(self, tmpdir_fn, model_cls): onnx_pred = self._helper_onnx_inference( model=model, onnx_filename=onnx_filename, - series=self.ts_tg_with_static - if model.uses_static_covariates - else self.ts_tg, - past_covariates=self.ts_pc if model.uses_past_covariates else None, - future_covariates=self.ts_fc if model.uses_future_covariates else None, + series=self.ts_tg_with_static, + past_covariates=self.ts_pc, + future_covariates=self.ts_fc, )[0][0] # check that the predictions are similar @@ -140,13 +138,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): onnx_pred = self._helper_onnx_inference( model=model_loaded, onnx_filename=onnx_filename, - series=self.ts_tg_with_static - if model_loaded.uses_static_covariates - else self.ts_tg, - past_covariates=self.ts_pc if model_loaded.uses_past_covariates else None, - future_covariates=self.ts_fc - if model_loaded.uses_future_covariates - else None, + series=self.ts_tg_with_static, + past_covariates=self.ts_pc, + future_covariates=self.ts_fc, )[0][0] # check that the predictions are similar @@ -176,13 +170,9 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): onnx_pred_weights = self._helper_onnx_inference( model=model_weights, onnx_filename=onnx_filename2, - series=self.ts_tg_with_static - if model_weights.uses_static_covariates - else self.ts_tg, - past_covariates=self.ts_pc if model_weights.uses_past_covariates else None, - future_covariates=self.ts_fc - if model_weights.uses_future_covariates - else None, + series=self.ts_tg_with_static, + past_covariates=self.ts_pc, + future_covariates=self.ts_fc, )[0][0] assert pred_weights.shape == onnx_pred_weights.shape, ( From dea831419f771d90b9bc7b2591b8a01640e66c8b Mon Sep 17 00:00:00 2001 From: madtoinou Date: Fri, 7 Mar 2025 17:25:50 +0100 Subject: [PATCH 16/17] address review comments --- .github/workflows/develop.yml | 8 +-- .github/workflows/merge.yml | 6 +-- .../forecasting/torch_forecasting_model.py | 4 +- darts/tests/conftest.py | 9 ++++ .../test_torch_forecasting_model.py | 4 +- darts/tests/optional_deps/test_onnx.py | 17 ++---- darts/tests/optional_deps/test_optuna.py | 6 ++- docs/userguide/hyperparameter_optimization.md | 10 ++-- docs/userguide/torch_forecasting_models.md | 54 +------------------ 9 files changed, 36 insertions(+), 82 deletions(-) diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml index 27773cfcd1..856a6edce7 100644 --- a/.github/workflows/develop.yml +++ b/.github/workflows/develop.yml @@ -99,7 +99,7 @@ jobs: - name: "Compile Dependency Versions" run: | curl -LsSf https://astral.sh/uv/install.sh | sh - uv pip compile requirements/dev-all.txt > requirements-latest.txt + uv pip compile requirements/dev-all.txt requirements/optional.txt > requirements-latest.txt # only restore cache but do not upload - name: "Restore cached python environment" @@ -120,7 +120,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt + pip install -U -r requirements/dev-all.txt -r requirements/optional.txt - name: "Install libomp (for LightGBM)" run: | @@ -152,7 +152,7 @@ jobs: - name: "Compile Dependency Versions" run: | curl -LsSf https://astral.sh/uv/install.sh | sh - uv pip compile requirements/dev-all.txt > requirements-latest.txt + uv pip compile requirements/dev-all.txt requirements/optional.txt > requirements-latest.txt # only restore cache but do not upload - name: "Restore cached python environment" @@ -169,7 +169,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt + pip install -U -r requirements/dev-all.txt -r requirements/optional.txt - name: "Install libomp (for LightGBM)" run: | diff --git a/.github/workflows/merge.yml b/.github/workflows/merge.yml index 22bccb81f4..b57ab115f4 100644 --- a/.github/workflows/merge.yml +++ b/.github/workflows/merge.yml @@ -111,7 +111,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt + pip install -U -r requirements/dev-all.txt -r requirements/optional.txt - name: "Install libomp (for LightGBM)" run: | @@ -141,7 +141,7 @@ jobs: - name: "Compile Dependency Versions" run: | curl -LsSf https://astral.sh/uv/install.sh | sh - uv pip compile requirements/dev-all.txt -o requirements-latest.txt + uv pip compile requirements/dev-all.txt requirements/optional.txt -o requirements-latest.txt # only restore cache but do not upload - name: "Restore cached python environment" @@ -162,7 +162,7 @@ jobs: - name: "Install Dependencies" run: | # install latest dependencies (potentially updating cached dependencies) - pip install -U -r requirements/dev-all.txt + pip install -U -r requirements/dev-all.txt -r requirements/optional.txt - name: "Install libomp (for LightGBM)" run: | diff --git a/darts/models/forecasting/torch_forecasting_model.py b/darts/models/forecasting/torch_forecasting_model.py index 3a003bcf25..c1e4b80dee 100644 --- a/darts/models/forecasting/torch_forecasting_model.py +++ b/darts/models/forecasting/torch_forecasting_model.py @@ -1782,7 +1782,7 @@ def save( self.trainer.save_checkpoint(path_ptl_ckpt, weights_only=clean) # TODO: keep track of PyTorch Lightning to see if they implement model checkpoint saving - # without having to call fit/predict/validate/test before + # without having to call fit/predict/validate/test before # try to recover original automatic PL checkpoint elif self.load_ckpt_path: if os.path.exists(self.load_ckpt_path): @@ -2141,7 +2141,7 @@ def load_weights_from_checkpoint( self.model.load_state_dict(ckpt["state_dict"], strict=strict) # update the fit_called attribute to allow for direct inference self._fit_called = True - # based on the shape of train_sample, figure out with covariates are used by the model + # based on the shape of train_sample, figure out which covariates are used by the model # (usually set in the Darts model prior to fitting it) self._update_covariates_use() diff --git a/darts/tests/conftest.py b/darts/tests/conftest.py index 3c7288b525..48ba0c5e25 100644 --- a/darts/tests/conftest.py +++ b/darts/tests/conftest.py @@ -50,6 +50,15 @@ } } +tfm_kwargs_dev = { + "pl_trainer_kwargs": { + "accelerator": "cpu", + "enable_progress_bar": False, + "enable_model_summary": False, + "fast_dev_run": True, + } +} + @pytest.fixture(scope="session", autouse=True) def set_up_tests(request): diff --git a/darts/tests/models/forecasting/test_torch_forecasting_model.py b/darts/tests/models/forecasting/test_torch_forecasting_model.py index 33557f37cb..0e9812761b 100644 --- a/darts/tests/models/forecasting/test_torch_forecasting_model.py +++ b/darts/tests/models/forecasting/test_torch_forecasting_model.py @@ -13,7 +13,7 @@ from darts.dataprocessing.encoders import SequentialEncoder from darts.dataprocessing.transformers import BoxCox, Scaler from darts.metrics import mape -from darts.tests.conftest import TORCH_AVAILABLE, tfm_kwargs +from darts.tests.conftest import TORCH_AVAILABLE, tfm_kwargs, tfm_kwargs_dev if not TORCH_AVAILABLE: pytest.skip( @@ -447,7 +447,7 @@ def test_save_and_load_weights_covs_usage_attributes(self, tmpdir_fn, params): input_chunk_length=4, output_chunk_length=1, n_epochs=1, - **tfm_kwargs, + **tfm_kwargs_dev, ) # skip test if the combination of covariates is not supported by the model if ( diff --git a/darts/tests/optional_deps/test_onnx.py b/darts/tests/optional_deps/test_onnx.py index 55b7d02091..1d31d27d20 100644 --- a/darts/tests/optional_deps/test_onnx.py +++ b/darts/tests/optional_deps/test_onnx.py @@ -7,7 +7,7 @@ import darts.utils.timeseries_generation as tg from darts import TimeSeries -from darts.tests.conftest import ONNX_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs +from darts.tests.conftest import ONNX_AVAILABLE, TORCH_AVAILABLE, tfm_kwargs_dev from darts.utils.onnx_utils import prepare_onnx_inputs if not (TORCH_AVAILABLE and ONNX_AVAILABLE): @@ -20,24 +20,15 @@ from darts.models import ( BlockRNNModel, - NBEATSModel, NHiTSModel, - NLinearModel, TiDEModel, - TransformerModel, ) # TODO: check how RINorm can be handled with respect to ONNX torch_model_cls = [ BlockRNNModel, - NBEATSModel, NHiTSModel, - NLinearModel, - # TCNModel, - # TFTModel, TiDEModel, - TransformerModel, - # TSMixerModel, ] @@ -54,7 +45,7 @@ class TestOnnx: @pytest.mark.parametrize("model_cls", torch_model_cls) def test_onnx_save_load(self, tmpdir_fn, model_cls): model = model_cls( - input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs + input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs_dev ) onnx_filename = f"test_onnx_{model.model_name}.onnx" @@ -103,7 +94,7 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): """Check that creating the onnx export from a model directly loaded from a checkpoint work as expected""" model_cls, clean = params model = model_cls( - input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs + input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs_dev ) onnx_filename = f"test_onnx_{model.model_name}.onnx" onnx_filename2 = f"test_onnx_{model.model_name}_weights.onnx" @@ -149,7 +140,7 @@ def test_onnx_from_ckpt(self, tmpdir_fn, params): # load only the weights model_weights = model_cls( - input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs + input_chunk_length=4, output_chunk_length=2, n_epochs=1, **tfm_kwargs_dev ) model_weights.load_weights(ckpt_filename) pred_weights = model_weights.predict( diff --git a/darts/tests/optional_deps/test_optuna.py b/darts/tests/optional_deps/test_optuna.py index 41368ba793..b49bdec7f2 100644 --- a/darts/tests/optional_deps/test_optuna.py +++ b/darts/tests/optional_deps/test_optuna.py @@ -25,7 +25,9 @@ # hacky workaround found in https://github.com/Lightning-AI/pytorch-lightning/issues/17485 # to avoid import of both lightning and pytorch_lightning - class PatchedCallback(optuna.integration.PyTorchLightningPruningCallback, Callback): + class PatchedPruningCallback( + optuna.integration.PyTorchLightningPruningCallback, Callback + ): pass from darts.models import TCNModel @@ -60,7 +62,7 @@ def objective(trial): include_year = trial.suggest_categorical("year", [False, True]) # throughout training we'll monitor the validation loss for both pruning and early stopping - pruner = PatchedCallback(trial, monitor="val_loss") + pruner = PatchedPruningCallback(trial, monitor="val_loss") early_stopper = EarlyStopping( "val_loss", min_delta=0.001, patience=3, verbose=True ) diff --git a/docs/userguide/hyperparameter_optimization.md b/docs/userguide/hyperparameter_optimization.md index 5097532424..8094b15698 100644 --- a/docs/userguide/hyperparameter_optimization.md +++ b/docs/userguide/hyperparameter_optimization.md @@ -20,7 +20,7 @@ import numpy as np import optuna import torch from optuna.integration import PyTorchLightningPruningCallback -from pytorch_lightning.callbacks import EarlyStopping +from pytorch_lightning.callbacks import Callback, EarlyStopping from sklearn.preprocessing import MaxAbsScaler from darts.dataprocessing.transformers import Scaler @@ -41,6 +41,11 @@ scaler = Scaler(MaxAbsScaler()) train = scaler.fit_transform(train) val = scaler.transform(val) +# workaround found in https://github.com/Lightning-AI/pytorch-lightning/issues/17485 +# to avoid import of both lightning and pytorch_lightning +class PatchedPruningCallback(optuna.integration.PyTorchLightningPruningCallback, Callback): + pass + # define objective function def objective(trial): # select input and output chunk lengths @@ -57,7 +62,7 @@ def objective(trial): include_year = trial.suggest_categorical("year", [False, True]) # throughout training we'll monitor the validation loss for both pruning and early stopping - pruner = PyTorchLightningPruningCallback(trial, monitor="val_loss") + pruner = PatchedPruningCallback(trial, monitor="val_loss") early_stopper = EarlyStopping("val_loss", min_delta=0.001, patience=3, verbose=True) callbacks = [pruner, early_stopper] @@ -112,7 +117,6 @@ def objective(trial): model.fit( series=train, val_series=model_val_set, - num_loader_workers=num_workers, ) # reload best model over course of training diff --git a/docs/userguide/torch_forecasting_models.md b/docs/userguide/torch_forecasting_models.md index 8b65345bd0..4aed55cf55 100644 --- a/docs/userguide/torch_forecasting_models.md +++ b/docs/userguide/torch_forecasting_models.md @@ -372,59 +372,7 @@ import onnx import onnxruntime as ort import numpy as np from darts import TimeSeries - -# can be imported from darts.utils.onnx_utils.py -def prepare_onnx_inputs( - model, - series: TimeSeries, - past_covariates : Optional[TimeSeries] = None, - future_covariates : Optional[TimeSeries] = None, -) -> tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]: - """Helper function to slice and concatenate the input features. - - In order to remove the dependency on the `model` argument, it can be decomposed into - the following arguments (and simplified depending on the characteristics of the model used): - - model_icl - - model_ocl - - model_uses_past_covs - - model_uses_future_covs - - model_uses_static_covs - """ - past_feats, future_feats, static_feats = None, None, None - # get input & output windows - past_start = series.end_time() - (model.input_chunk_length - 1) * series.freq - past_end = series.end_time() - future_start = past_end + 1 * series.freq - future_end = past_end + model.output_chunk_length * series.freq - # extract all historic and future features from target, past and future covariates - past_feats = series[past_start:past_end].values() - if past_covariates and model.uses_past_covariates: - # extract past covariates - past_feats = np.concatenate( - [ - past_feats, - past_covariates[past_start:past_end].values() - ], - axis=1 - ) - if future_covariates and model.uses_future_covariates: - # extract past part of future covariates - past_feats = np.concatenate( - [ - past_feats, - future_covariates[past_start:past_end].values() - ], - axis=1 - ) - # extract future part of future covariates - future_feats = future_covariates[future_start:future_end].values() - # add batch dimension -> (batch, n time steps, n components) - past_feats = np.expand_dims(past_feats, axis=0).astype(series.dtype) - future_feats = np.expand_dims(future_feats, axis=0).astype(series.dtype) - # extract static covariates - if series.has_static_covariates and model.uses_static_covariates: - static_feats = np.expand_dims(series.static_covariates_values(), axis=0).astype(series.dtype) - return past_feats, future_feats, static_feats +from darts.utils.onnx_utils.py import prepare_onnx_inputs onnx_model = onnx.load(onnx_filename) onnx.checker.check_model(onnx_model) From 2a7c268db1d730662f7433e71aed1e62c991c1de Mon Sep 17 00:00:00 2001 From: dennisbader Date: Fri, 7 Mar 2025 17:38:14 +0100 Subject: [PATCH 17/17] minor update --- darts/tests/models/forecasting/test_torch_forecasting_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/darts/tests/models/forecasting/test_torch_forecasting_model.py b/darts/tests/models/forecasting/test_torch_forecasting_model.py index 0e9812761b..37997f9d0a 100644 --- a/darts/tests/models/forecasting/test_torch_forecasting_model.py +++ b/darts/tests/models/forecasting/test_torch_forecasting_model.py @@ -470,7 +470,7 @@ def test_save_and_load_weights_covs_usage_attributes(self, tmpdir_fn, params): model_loaded = model_cls( input_chunk_length=4, output_chunk_length=1, - **tfm_kwargs, + **tfm_kwargs_dev, ) model_loaded.load_weights(filename_ckpt)