Auto batch size for torch model #2318

BohdanBilonoh · 2024-04-11T19:57:11Z

Checklist before merging this PR:

Mentioned all issues that this PR fixes or addresses.
Summarized the updates of this PR under Summary.
Added an entry under Unreleased in the Changelog.

Summary

Auto batch size finding for TorchForecastingModel. This is just a wrapper for lightning.pytorch.tuner.tuning.Tuner.scale_batch_size

codecov · 2024-04-11T20:33:46Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.02%. Comparing base (a0cc279) to head (35d1d58).

❗ Current head 35d1d58 differs from pull request most recent head ae7a128

Please upload reports for the commit ae7a128 to get more accurate results.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2318      +/-   ##
==========================================
+ Coverage   93.75%   94.02%   +0.26%     
==========================================
  Files         138      138              
  Lines       14352    14152     -200     
==========================================
- Hits        13456    13306     -150     
+ Misses        896      846      -50

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

VascoSch92 · 2024-04-12T06:27:03Z

darts/models/forecasting/torch_forecasting_model.py

+        max_samples_per_ts: Optional[int] = None,
+        num_loader_workers: int = 0,
+        method: Literal["fit", "validate", "test", "predict"] = "fit",
+        mode: str = "power",


Why are you not using Literal here? This variable can just be power or linear right?

I took it from lightning. I think this is motivated by fact that the mode potentially could be many more modes

dennisbader

Thanks for implementing this @BohdanBilonoh. It looks already pretty good 🚀
I added some suggestions on how we could simplify things a bit, and also to support "predict" method.

dennisbader · 2024-04-12T11:58:16Z

darts/models/forecasting/torch_forecasting_model.py

+        epochs: int = 0,
+        max_samples_per_ts: Optional[int] = None,
+        num_loader_workers: int = 0,


These are not required I guess

Suggested change

epochs: int = 0,

max_samples_per_ts: Optional[int] = None,

num_loader_workers: int = 0,

dennisbader · 2024-04-12T12:00:31Z

darts/models/forecasting/torch_forecasting_model.py

+        epochs: int = 0,
+        max_samples_per_ts: Optional[int] = None,
+        num_loader_workers: int = 0,
+        method: Literal["fit", "validate", "test", "predict"] = "fit",


"test" is not supported for darts.
"predict" would require a datamodule for prediction

in my opinion, we should only support "fit" and "predict", since we use the same batch size for train and val.

dennisbader · 2024-04-12T12:07:24Z

darts/models/forecasting/torch_forecasting_model.py

+        batch_arg_name
+            The name of the argument to scale in the model. Defaults to 'batch_size'.


not required

Suggested change

batch_arg_name

The name of the argument to scale in the model. Defaults to 'batch_size'.

dennisbader · 2024-04-12T12:10:28Z

darts/models/forecasting/torch_forecasting_model.py

+                self.batch_size = batch_size
+
+            def train_dataloader(self):
+                return DataLoader(


since we use this also in _setup_for_train and predict_from_dataset, it would be good to have this logic in a private method for example _build_dataloader(), that takes as input a dataset, and returns the dataloader according to "train", "val" and "predict"

dennisbader · 2024-04-12T12:14:54Z

darts/models/forecasting/torch_forecasting_model.py

+    def scale_batch_size(
+        self,
+        series: Union[TimeSeries, Sequence[TimeSeries]],
+        val_series: Union[TimeSeries, Sequence[TimeSeries]],


I think we can remove all val_* arguments, since we use the same batch size for training and evaluation.

For the val_dataloader in the DataModule, we can just use series, past_covariates, future_covariates for the input dataset

dennisbader · 2024-04-12T12:15:22Z

darts/tests/models/forecasting/test_torch_forecasting_model.py

@@ -1373,6 +1373,21 @@ def test_lr_find(self):
            )
        assert scores["worst"] > scores["suggested"]

+    @pytest.mark.slow


it's not slow

Suggested change

@pytest.mark.slow

dennisbader · 2024-04-12T12:16:01Z

darts/tests/models/forecasting/test_torch_forecasting_model.py

@@ -1373,6 +1373,21 @@ def test_lr_find(self):
            )
        assert scores["worst"] > scores["suggested"]

+    @pytest.mark.slow
+    def test_scale_batch_size(self):


test this for method="fit" and "predict"

dennisbader · 2024-04-12T12:17:32Z

CHANGELOG.md

@@ -89,6 +89,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
  - Renamed the private `_is_probabilistic` property to a public `supports_probabilistic_prediction`.
 - Improvements to `DataTransformer`: [#2267](https://github.com/unit8co/darts/pull/2267) by [Alicja Krzeminska-Sciga](https://github.com/alicjakrzeminska).
  - `InvertibleDataTransformer` now supports parallelized inverse transformation for `series` being a list of lists of `TimeSeries` (`Sequence[Sequence[TimeSeries]]`). This `series` type represents for example the output from `historical_forecasts()` when using multiple series. 
+- New method `TorchForecastingModel.scale_batch_size()` that helps to find batch size automatically. [#2318](https://github.com/unit8co/darts/pull/2318) by [Bohdan Bilonoh](https://github.com/BohdanBilonoh)


Suggested change

- New method `TorchForecastingModel.scale_batch_size()` that helps to find batch size automatically. [#2318](https://github.com/unit8co/darts/pull/2318) by [Bohdan Bilonoh](https://github.com/BohdanBilonoh)

- Improvements to `TorchForecastingModel`:

- New method `TorchForecastingModel.scale_batch_size()` to find the maximum batch size for fit and predict before memory would run out. [#2318](https://github.com/unit8co/darts/pull/2318) by [Bohdan Bilonoh](https://github.com/BohdanBilonoh)

BohdanBilonoh · 2024-05-28T12:29:43Z

@dennisbader could you please help with predict mode. It is tricky because TorchForecastingModel requires set_predict_parameters to be called before predict. Unfortunately I don't know how to make it works

BohdanBilonoh requested review from dennisbader and madtoinou as code owners April 11, 2024 19:57

VascoSch92 reviewed Apr 12, 2024

View reviewed changes

dennisbader requested changes Apr 12, 2024

View reviewed changes

Bohdan Bilonoh added 3 commits May 28, 2024 12:02

Auto batch size for torch model

3667f9e

update CHANGELOG.md

46d8a42

remove code smell

932514e

WIP: stuck with model.set_predict_parameters

ae7a128

BohdanBilonoh force-pushed the feature/pl_auto_batch_size branch from eecfd0f to ae7a128 Compare May 28, 2024 12:30

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Auto batch size for torch model #2318

Auto batch size for torch model #2318

BohdanBilonoh commented Apr 11, 2024 •

edited

Loading

codecov bot commented Apr 11, 2024 •

edited

Loading

VascoSch92 Apr 12, 2024

BohdanBilonoh Apr 12, 2024 •

edited

Loading

dennisbader left a comment

dennisbader Apr 12, 2024

dennisbader Apr 12, 2024

dennisbader Apr 12, 2024

dennisbader Apr 12, 2024

dennisbader Apr 12, 2024

dennisbader Apr 12, 2024

dennisbader Apr 12, 2024

dennisbader Apr 12, 2024

BohdanBilonoh commented May 28, 2024

	epochs: int = 0,
	max_samples_per_ts: Optional[int] = None,
	num_loader_workers: int = 0,

		batch_arg_name
		The name of the argument to scale in the model. Defaults to 'batch_size'.

	- New method `TorchForecastingModel.scale_batch_size()` that helps to find batch size automatically. [#2318](https://github.com/unit8co/darts/pull/2318) by [Bohdan Bilonoh](https://github.com/BohdanBilonoh)
	- Improvements to `TorchForecastingModel`:
	- New method `TorchForecastingModel.scale_batch_size()` to find the maximum batch size for fit and predict before memory would run out. [#2318](https://github.com/unit8co/darts/pull/2318) by [Bohdan Bilonoh](https://github.com/BohdanBilonoh)

Auto batch size for torch model #2318

Are you sure you want to change the base?

Auto batch size for torch model #2318

Conversation

BohdanBilonoh commented Apr 11, 2024 • edited Loading

Summary

codecov bot commented Apr 11, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

BohdanBilonoh Apr 12, 2024 • edited Loading

Choose a reason for hiding this comment

dennisbader left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BohdanBilonoh commented May 28, 2024

BohdanBilonoh commented Apr 11, 2024 •

edited

Loading

codecov bot commented Apr 11, 2024 •

edited

Loading

BohdanBilonoh Apr 12, 2024 •

edited

Loading