Feat/onnx support (#2620)

* feat: wrapping around pl.to_onnx to export models to ONNX, still require testing

* feat: cleaned implementation of the to_onnx method

* fix: generation of input name, shape of input_batch for PastCov torch module

* feat: adding example of onnx usage in userguide

* update changelog

* fix: revert some changes

* fix: export to onnx for RNNModel

* feat: added a comment about RNNModel for onnx inference

* update changelog

* fix: address review comments

* update changelog

* update torch user guide

* update to_onnx

---------

Co-authored-by: dennisbader <[email protected]>

Author: madtoinou

CHANGELOG.md

@@ -11,6 +11,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Improved**
 
+- Added ONNX support for torch-based models with method `TorchForecastingModel.to_onnx()`. Check out [this example](https://unit8co.github.io/darts/userguide/gpu_and_tpu_usage.html#exporting-model-to-onnx-format-for-inference) from the user guide on how to export and load a model for inference. [#2620](https://github.com/unit8co/darts/pull/2620) by [Antoine Madrona](https://github.com/madtoinou)
 - Made method `ForecastingModel.untrained_model()` public. Use this method to get a new (untrained) model instance created with the same parameters. [#2684](https://github.com/unit8co/darts/pull/2684) by [Timon Erhart](https://github.com/turbotimon)
 
 **Fixed**

darts/models/forecasting/pl_forecasting_module.py

@@ -93,6 +93,7 @@ def __init__(
         When subclassing this class, please make sure to add the following methods with the given signatures:
             - :func:`PLForecastingModule.__init__()`
             - :func:`PLForecastingModule.forward()`
+            - :func:`PLForecastingModule._process_input_batch()`
             - :func:`PLForecastingModule._produce_train_output()`
             - :func:`PLForecastingModule._get_batch_prediction()`
 
@@ -632,17 +633,48 @@ def _produce_train_output(self, input_batch: tuple):
         input_batch
             ``(past_target, past_covariates, static_covariates)``
         """
-        past_target, past_covariates, static_covariates = input_batch
+        return self(self._process_input_batch(input_batch))
+
+    def _process_input_batch(
+        self, input_batch: tuple
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Converts output of PastCovariatesDataset (training dataset) into an input/past- and
+        output/future chunk.
+
+        Parameters
+        ----------
+        input_batch
+            ``(past_target, past_covariates, historic_future_covariates, future_covariates, static_covariates)``.
+
+        Returns
+        -------
+        tuple
+            ``(x_past, x_static)`` the input/past and output/future chunks.
+        """
+        # because of future past covariates, the batch shape is different during training and prediction
+        if len(input_batch) == 3:
+            (
+                past_target,
+                past_covariates,
+                static_covariates,
+            ) = input_batch
+        else:
+            (
+                past_target,
+                past_covariates,
+                future_past_covariates,
+                static_covariates,
+            ) = input_batch
         # Currently all our PastCovariates models require past target and covariates concatenated
-        inpt = (
+        return (
             (
                 torch.cat([past_target, past_covariates], dim=2)
                 if past_covariates is not None
                 else past_target
             ),
             static_covariates,
         )
-        return self(inpt)
 
     def _get_batch_prediction(
         self, n: int, input_batch: tuple, roll_size: int
@@ -674,12 +706,9 @@ def _get_batch_prediction(
             past_covariates.shape[dim_component] if past_covariates is not None else 0
         )
 
-        input_past = torch.cat(
-            [ds for ds in [past_target, past_covariates] if ds is not None],
-            dim=dim_component,
-        )
+        input_past, input_static = self._process_input_batch(input_batch)
 
-        out = self._produce_predict_output(x=(input_past, static_covariates))[
+        out = self._produce_predict_output(x=(input_past, input_static))[
             :, self.first_prediction_index :, :
         ]
 
@@ -796,7 +825,7 @@ def _process_input_batch(
             future_covariates,
             static_covariates,
         ) = input_batch
-        dim_variable = 2
+        dim_comp = 2
 
         x_past = torch.cat(
             [
@@ -808,7 +837,7 @@ def _process_input_batch(
                 ]
                 if tensor is not None
             ],
-            dim=dim_variable,
+            dim=dim_comp,
         )
         return x_past, future_covariates, static_covariates
 

darts/models/forecasting/rnn_model.py

@@ -104,6 +104,12 @@ def forward(
         pass
 
     def _produce_train_output(self, input_batch: tuple) -> torch.Tensor:
+        # only return the forecast, not the hidden state
+        return self(self._process_input_batch(input_batch))[0]
+
+    def _process_input_batch(
+        self, input_batch: tuple
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         (
             past_target,
             historic_future_covariates,
@@ -112,15 +118,14 @@ def _produce_train_output(self, input_batch: tuple) -> torch.Tensor:
         ) = input_batch
         # For the RNN we concatenate the past_target with the future_covariates
         # (they have the same length because we enforce a Shift dataset for RNNs)
-        model_input = (
+        return (
             (
                 torch.cat([past_target, future_covariates], dim=2)
                 if future_covariates is not None
                 else past_target
             ),
             static_covariates,
         )
-        return self(model_input)[0]
 
     def _produce_predict_output(
         self, x: tuple, last_hidden_state: Optional[torch.Tensor] = None

darts/models/forecasting/torch_forecasting_model.py

@@ -646,6 +646,66 @@ def _verify_past_future_covariates(self, past_covariates, future_covariates):
                 logger=logger,
             )
 
+    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 <https://lightning.ai/docs/pytorch/
+        stable/common/lightning_module.html#to-onnx>`_).
+
+        Note: requires `onnx` library (optional dependency) to be installed.
+
+        Example for exporting a :class:`DLinearModel`:
+
+        .. highlight:: python
+        .. code-block:: python
+
+            from darts.datasets import AirPassengersDataset
+            from darts.models import DLinearModel
+
+            series = AirPassengersDataset().load()
+            model = DLinearModel(input_chunk_length=4, output_chunk_length=1)
+            model.fit(series, epochs=1)
+            model.to_onnx("my_model.onnx")
+        ..
+
+        Parameters
+        ----------
+        path
+            Path under which to save the model at its current state. If no path is specified, the model
+            is automatically saved under ``"{ModelClass}_{YYYY-mm-dd_HH_MM_SS}.onnx"``.
+        **kwargs
+            Additional kwargs for PyTorch's :func:`torch.onnx.export` method (except parameters ``file_path``,
+            ``input_sample``, ``input_name``). For more information, read the `official documentation
+            <https://pytorch.org/docs/master/onnx.html#torch.onnx.export>`_.
+        """
+        if not self._fit_called:
+            raise_log(
+                ValueError("`fit()` needs to be called before `to_onnx()`."), logger
+            )
+
+        if path is None:
+            path = self._default_save_path() + ".onnx"
+
+        # last dimension in train_sample_shape is the expected target
+        mock_batch = tuple(
+            torch.rand((1,) + shape, dtype=self.model.dtype) if shape else None
+            for shape in self.model.train_sample_shape[:-1]
+        )
+        input_sample = self.model._process_input_batch(mock_batch)
+
+        # torch models necessarily use historic target values as features in current implementation
+        input_names = ["x_past"]
+        if self.uses_future_covariates:
+            input_names.append("x_future")
+        if self.uses_static_covariates:
+            input_names.append("x_static")
+
+        self.model.to_onnx(
+            file_path=path,
+            input_sample=(input_sample,),
+            input_names=input_names,
+            **kwargs,
+        )
+
     @random_method
     def fit(
         self,

docs/userguide/torch_forecasting_models.md

@@ -22,6 +22,7 @@ We assume that you already know about covariates in Darts. If you're new to the
       - [Manual saving / loading](#manual-saving--loading)
       - [Train & save on GPU, load on CPU](#trainingsaving-on-gpu-and-loading-on-cpu)
       - [Load pre-trained model for fine-tuning](#re-training-or-fine-tuning-a-pre-trained-model)
+      - [Exporting model to ONNX format for inference](#exporting-model-to-ONNX-format-for-inference)
     - [Callbacks](#callbacks)
       - [Early Stopping](#example-with-early-stopping)
       - [Custom Callback](#example-of-custom-callback-to-store-losses)
@@ -350,6 +351,93 @@ model_finetune = SomeTorchForecastingModel(...,  # use identical parameters & va
 model_finetune.load_weights("/your/path/to/save/model.pt")
 ```
 
+#### Exporting model to ONNX format for inference
+
+It is also possible to export the model weights to the ONNX format to run inference in a lightweight environment. The example below works for any `TorchForecastingModel` except `RNNModel` and for optional usage of past, future and / or static covariates. Note that all series and covariates must extend far enough into the past (`input_chunk_length)` and future (`output_chunk_length`) relative to the end of the target `series`. It will not be possible to forecast a horizon `n > output_chunk_length` without implementing the auto-regression logic.
+
+```python
+model = SomeTorchForecastingModel(...)
+model.fit(...)
+
+# make sure to have `onnx` and `onnxruntime` installed
+onnx_filename = "example_onnx.onnx"
+model.to_onnx(onnx_filename, export_params=True)
+```
+
+Now, to load the model and predict steps after the end of the series:
+
+```python
+from typing import Optional
+import onnx
+import onnxruntime as ort
+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[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
+
+onnx_model = onnx.load(onnx_filename)
+onnx.checker.check_model(onnx_model)
+ort_session = ort.InferenceSession(onnx_filename)
+
+# use helper function to extract the features from the series
+past_feats, future_feats, static_feats = prepare_onnx_inputs(
+    model=model,
+    series=series,
+    past_covariates=ts_past,
+    future_covariates=ts_future,
+)
+
+# 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)
+ort_out = ort_session.run(None, ort_inputs)
+```
+
 ### Callbacks
 
 Callbacks are a powerful way to monitor or control the behavior of the model during the training process. Some examples: