[ENH] Implement a unified hyperparameter tuning interface decoupled from model classes for v2 by Muhammad-Rebaal · Pull Request #2335 · sktime/pytorch-forecasting

Muhammad-Rebaal · 2026-07-07T14:56:39Z

Fixes : #2332

This PR implements the hyperparameter optimization interface for v2 models.

Implemented Architecture

SearchRange & Global Registry (tuning/search_range.py, tuning/global_registry.py):
- A typed container for search spaces and a global dictionary defining standard ranges for common parameters (e.g., hidden_size, dropout, n_heads).
Auto-Discovery Hook (models/base/_base_model_v2.py):
- A get_tuneable_hyperparameters() class method on BaseModel that uses inspect.signature(cls.__init__) to detect what parameters the subclass accepts and maps them to the global registry.
HyperparameterTuner (tuning/hyperparameter_tuner.py):
- A standalone optimizer living outside model modules. It queries the base hook, merges user overrides, and runs sequential Optuna trials using Base_pkg.fit().

User Flow:

from pytorch_forecasting.models.timexer import TimeXer_pkg_v2
from pytorch_forecasting.models.dlinear import DLinear_pkg_v2
from pytorch_forecasting.tuning.hyperparameter_tuner import HyperparameterTuner
from pytorch_forecasting.tuning.search_range import SearchRange

tuner = HyperparameterTuner(
    pkg_cls=TimeXer_pkg_v2,
    data=my_timeseries_data,
    base_datamodule_cfg={"context_length": 96, "prediction_length": 24},
)
study = tuner.optimize(n_trials=50, max_epochs=10)
print(f"Best params: {study.best_params}")
print(f"Best val_loss: {study.best_value}")

…from model classes for v2

codecov · 2026-07-07T16:21:38Z

Codecov Report

❌ Patch coverage is 86.27451% with 14 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@0acb5a0). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
pytorch_forecasting/tuning/hyperparameter_tuner.py	86.00%	14 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2335   +/-   ##
=======================================
  Coverage        ?   88.85%           
=======================================
  Files           ?      219           
  Lines           ?    11772           
  Branches        ?        0           
=======================================
  Hits            ?    10460           
  Misses          ?     1312           
  Partials        ?        0

Flag	Coverage Δ
cpu	`88.85% <86.27%> (?)`
pytest	`88.85% <86.27%> (?)`

Flags with carried forward coverage won't be shown. Click here to find out more.

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

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

phoeenniixx

Why aren't we using optuna here? Also, in the param grid the values are hard coded

Muhammad-Rebaal · 2026-07-08T14:05:37Z

Why aren't we using optuna here?

Because of these 2 reason :

User Experience: It's much easier for a user to inspect, modify, and pass a simple SearchRange dataclass in custom_ranges than it is to write and pass custom Optuna lambda functions.

Decoupling: By keeping Optuna logic out of the base model layer and registry, the models themselves remain completely framework-agnostic. The SearchRange dataclass acts as a clean bridge.

param grid the values are hard coded

In v1's optimize_hyperparameters for TFT, the default search spaces were hardcoded as tuple arguments (e.g., hidden_size_range=(16, 265), dropout_range=(0.1, 0.3)). The global_registry simply takes those v1 hardcoded defaults and standardizes them centrally so they can apply to all models automatically. However, they aren't strictly locked users can completely override these default "hardcoded" values by passing their own custom_ranges to the HyperparameterTuner.

fkiraly

Hm, ok start. I would make all classes and variables private for now, so we can later change. This increases separation of concerns which is good.

Some questions in the design:

how would this get used?
This designs custom search range classes. Do we want to do that, instead of relying on one of the standards? (e.g., even optuna itself)
optuna is one choice of what we could do. Have a look at the hyperactive tuning class for neural networks, or the OptCV class in sktime. Could this be useful?

Muhammad-Rebaal · 2026-07-12T22:05:19Z

Thanks for the review.

Hm, ok start. I would make all classes and variables private for now, so we can later change. This increases separation of concerns which is good.

I've implemented this.

how would this get used?

from pytorch_forecasting.tuning import ForecastingSearchCV, _SearchRange
from pytorch_forecasting.models.dlinear import DLinear_pkg_v2
from pytorch_forecasting.metrics import SMAPE


datasets = DLinear_pkg_v2.get_test_dataset_from()

search = ForecastingSearchCV(
    pkg_cls=DLinear_pkg_v2,
    param_grid={
        "moving_avg": _SearchRange(param_type="categorical", choices=[3, 5, 7, 11]),
        "dropout": _SearchRange(param_type="float", low=0.1, high=0.5)
    },
    n_trials=3,
    base_model_cfg={"loss": SMAPE()},  
    base_datamodule_cfg={
        "context_length": 8,        
        "prediction_length": 2,
        "batch_size": 2,                  
        "train_val_test_split": (0.5, 0.5)
    },  
    base_trainer_cfg={
        "max_epochs": 2,
        "limit_train_batches": 2,         
        "limit_val_batches": 1            
    },      
)

search.fit(datasets["train"])

predictions = search.predict(datasets["predict"])

print("Best parameters found:", search.best_params_)

This designs custom search range classes. Do we want to do that, instead of relying on one of the standards? (e.g., even optuna itself)

_SearchRange is intentionally a thin adapter that maps to Optuna internally see suggest() method. If we ever want to support a second backend, e.g., Ray Tune, we only change suggest() no user-facing code changes. Using raw optuna.distributions would hard-couple every model registry entry to Optuna.

optuna is one choice of what we could do. Have a look at the hyperactive tuning class for neural networks, or the OptCV class in sktime. Could this be useful?

Yes, it is preciously useful, helped me a lot and implemented the wrapper class ForecastingSearchCV like it is in sktime. Also with the help of the param_grid we don't have to create hooks in the base_model class as well.

phoeenniixx

Please use soft-dep isolation when importing soft dependencies
See https://www.sktime.net/docs/developer-guide/dependencies/#isolating-soft-dependencies-to-estimators
I think lightning.tuner can also be very helpful here, no?
Please add tests

phoeenniixx · 2026-07-13T08:22:49Z

+        try:
+            import optuna
+        except ImportError:
+            raise ImportError(
+                "Optuna is required for hyperparameter tuning. "


why aren't we using _safe_import or _check_soft_dependencies here?

phoeenniixx · 2026-07-13T08:28:00Z

+        custom_ranges : dict[str, SearchRange], optional
+            Override or add to auto-discovered ranges.
+        study : optuna.Study, optional
+            Existing study to resume.


what if there is no existing study and the user wants to create a new study here?

phoeenniixx · 2026-07-13T08:28:12Z

+        if custom_ranges:
+            search_ranges.update(custom_ranges)
+
+        def objective(trial):


should this be a private method?

I am talking about the method not the file?
And this is not even the file that I have added this comment on?

I just got a little confused, thanks. I'll make objective method private as its an internal one.

Please dont use AI blindly, this creates a lot of confusion and wastes time for all of us

phoeenniixx · 2026-07-13T08:30:33Z

from pytorch_forecasting.tuning import ForecastingSearchCV, _SearchRange
from pytorch_forecasting.models.dlinear import DLinear_pkg_v2
from pytorch_forecasting.metrics import SMAPE

datasets = DLinear_pkg_v2.get_test_dataset_from()

search = ForecastingSearchCV(
pkg_cls=DLinear_pkg_v2,
param_grid={
"moving_avg": _SearchRange(param_type="categorical", choices=[3, 5, 7, 11]),
"dropout": _SearchRange(param_type="float", low=0.1, high=0.5)
},
n_trials=3,
base_model_cfg={"loss": SMAPE()},
base_datamodule_cfg={
"context_length": 8,
"prediction_length": 2,
"batch_size": 2,
"train_val_test_split": (0.5, 0.5)
},
base_trainer_cfg={
"max_epochs": 2,
"limit_train_batches": 2,
"limit_val_batches": 1
},
)

search.fit(datasets["train"])

predictions = search.predict(datasets["predict"])

print("Best parameters found:", search.best_params_)

based on this vignette:
Why the user has to pass a private _SearchRange class? Why dont we create this inside the wrapper itself and the user passes the params more organically?

Also, why do we need to pass the cfgs again to this class, I thought we could simply pass the pkg class and it would handle the params itself?

Muhammad-Rebaal · 2026-07-14T14:55:24Z

from pytorch_forecasting.tuning import ForecastingSearchCV, _SearchRange
from pytorch_forecasting.models.dlinear import DLinear_pkg_v2
from pytorch_forecasting.metrics import SMAPE
datasets = DLinear_pkg_v2.get_test_dataset_from()
search = ForecastingSearchCV(
pkg_cls=DLinear_pkg_v2,
param_grid={
"moving_avg": _SearchRange(param_type="categorical", choices=[3, 5, 7, 11]),
"dropout": SearchRange(param_type="float", low=0.1, high=0.5)
},
n_trials=3,
base_model_cfg={"loss": SMAPE()},
base_datamodule_cfg={
"context_length": 8,
"prediction_length": 2,
"batch_size": 2,
"train_val_test_split": (0.5, 0.5)
},
base_trainer_cfg={
"max_epochs": 2,
"limit_train_batches": 2,
"limit_val_batches": 1
},
)
search.fit(datasets["train"])
predictions = search.predict(datasets["predict"])
print("Best parameters found:", search.best_params)

based on this vignette: Why the user has to pass a private _SearchRange class? Why dont we create this inside the wrapper itself and the user passes the params more organically?

Also, why do we need to pass the cfgs again to this class, I thought we could simply pass the pkg class and it would handle the params itself?

Nice Idea I'll try to implement that.

Muhammad-Rebaal · 2026-07-15T14:57:34Z

I think lightning.tuner can also be very helpful here, no?

Yes,
What I found is :

Tuner.lr_find() —> finds optimal learning rate
Tuner.scale_batch_size() —> finds max batch size that fits in memory
Ref : https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.tuner.tuning.Tuner.html

These are the methods I found that can be helpful. Any suggestions from your side.. @phoeenniixx

…dded code refactored and removed unwanted tests

…orch-forecasting into optimization

Muhammad-Rebaal · 2026-08-26T16:00:30Z

Hi @phoeenniixx,

Now the Hyperparameters can be tuned for any v2 model like this :

tuner = HyperparameterTuner(
    model="TFT",
    data=my_datamodule, 
)

Although the user can provide their trainer_cfg, and datamodule_cfg if passed the timeseries data.

…d trainer_cfg

…orch-forecasting into optimization

Muhammad-Rebaal · 2026-08-31T14:08:06Z

Hi @phoeenniixx,
I've made the following changes:

Removed datamodule_cfg, trainer_cfg.

The user can now pass raw timeseries and datamodule as well if passed raw timeseries datamodule would be made to processed and if lightening datamodule is passed then the user can also specify the train-test-split that we are working on another PR.

Lastly, every model has a get_datamodule_cls method implemented in their respective pkg coming from base_pkg, and since we are moving towards deprecation of the base_pkg. I've moved the get_datamodule_cls in their respective model_cls such as

BaseModel has EncoderDecoderTimeSeriesDataModule
TslibBaseModel override the BaseModel get_datamodule_cls method TslibDataModule

The purpose of this is that the user don't have to specify which datamodule is for which model and converts their data into the specific datamodule to be processed also added the edge case if the user passes the wrong datamodule it would raise a TypeError

Muhammad-Rebaal · 2026-09-10T13:36:24Z

Hi @phoeenniixx,
I've made the following change:

Added a variable datamodule_cls, now the user have to define the datamodule he has to use for the model

phoeenniixx · 2026-09-10T14:28:08Z

+            or against which to validate when ``data`` is a prebuilt DataModule.
+        **fixed_hparams
+            Any model parameter that should stay constant, e.g.
+            ``hidden_size=128``.


please add some examples here to how to use this

phoeenniixx

Thanks
Can you please clean up the test file? it has a lot of "param checks" like if i pass a, it is collected in x.param etc. It checks if something is implemented or not - we should only test if the things are working correctly... Not that if the tests accpet some param or not, it is understood that the class will accept datamodule and assign it to tuner.datamoudle.
(There are multiple tests like that here)

Also, pls add some examples (maybe even an example nb!)

…orch-forecasting into optimization

Muhammad-Rebaal · 2026-09-17T10:14:42Z

Hi @phoeenniixx,
I've cleaned up the test file and also added an example.

phoeenniixx

Thanks!
Will this work for a model checkpoint etc as well?
Also, should we also add a nb etc - as rn we cant verify that the complete workflow works as expected

phoeenniixx · 2026-09-19T20:01:25Z

+        if datamodule_cls is not None:
+            if not (
+                isinstance(datamodule_cls, type)
+                and issubclass(datamodule_cls, LightningDataModule)


do we need this validation? if we add type hints etc, it would be understood that it is a LightningDataModule no?

phoeenniixx · 2026-09-19T20:07:03Z

+            trainer.fit(model, datamodule=self.datamodule)
+
+            metrics = trainer.callback_metrics
+            for key in ("val_loss", "train_loss_epoch", "train_loss"):


should we optimize for train losses? that might lead to overfitting?

yes, you're right over here

phoeenniixx · 2026-09-19T20:08:05Z

+        custom_ranges = custom_ranges or {}
+
+        for param_name, param_obj in sig.parameters.items():
+            if param_name in base_params or param_name == "self":


Here what will happen to the params which are sent to theBaseModel? Like Samformer sends optimizer etc directly to base model:

pytorch-forecasting/pytorch_forecasting/models/samformer/_samformer_v2.py

Lines 61 to 69 in 0158d17

super().__init__(

loss=loss,

logging_metrics=logging_metrics,

optimizer=optimizer,

optimizer_params=optimizer_params,

lr_scheduler=lr_scheduler,

lr_scheduler_params=lr_scheduler_params,

)

Right now it gonna skip, in my mind we're tuning the model specific ones. I think skipping self and logging_metrics instead of whole base_model param would be a cleaner fix

but should we also skip optimizer, lr_scheduler etc?

For optimizer_params , lr_scheduler , these 2 have default none it would just auto skip as we've a logic:

if ( default is inspect.Parameter.empty or default is None or isinstance(default, (dict, list)) ): continue

For optimizer and loss, I think if we skip the users would lost the ability to pass their custom ranges ?

Optuna can't tune dictionaries like optimizer_params directly anyway. Since the learning rate is universally critical I think we just inject it directly inside the _objective function (just like the old TFT tuning.py did):

if "optimizer_params" not in self.fixed_hparams: lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True) model_cfg["optimizer_params"] = {"lr": lr}

I think i dont understand the issue - i think we need to optimise some of the params atleast - like lr_scheduler, no? So, how should we do that?

Can you please elaborate on this?

I thought above you're telling me not to tune lr_scheduler that's why I said we already have a logic to skip and suggested what we can do with other params.

My ques is that how are these params optimized here? Are we optimizing lr_scheduler? And come to think of it, for lr, we can also use tuner from lightning, no? That would be better?

Sorry, for not being clear :)
I just want to know how are the params that we are passing to base_model directly being optimized?

Muhammad-Rebaal added 2 commits July 3, 2026 19:58

feat : Implement a unified hyperparameter tuning interface decoupled …

bb7d41c

…from model classes for v2

feat : Implement a unified hyperparameter tuning interface decoupled …

0c691db

…from model classes for v2

Muhammad-Rebaal requested review from PranavBhatP, benHeid, fkiraly, fnhirwa, jdb78, phoeenniixx and yarnabrina as code owners July 7, 2026 14:56

fix pytest

772fa91

phoeenniixx requested changes Jul 8, 2026

View reviewed changes

fkiraly requested changes Jul 10, 2026

View reviewed changes

feat: Implemented Design wrapper similar to sktime

8301098

phoeenniixx requested changes Jul 13, 2026

View reviewed changes

phoeenniixx changed the title ~~Feat : Implement a unified hyperparameter tuning interface decoupled from model classes for v2~~ [ENH] Implement a unified hyperparameter tuning interface decoupled from model classes for v2 Jul 14, 2026

phoeenniixx added enhancement New feature or request module:models ptf-v2 Related to `pytorch-forecasting` v2 labels Jul 14, 2026

phoeenniixx assigned Muhammad-Rebaal Jul 14, 2026

Muhammad-Rebaal and others added 2 commits July 14, 2026 21:21

Merge branch 'main' into optimization

5133352

passed the cfgs to class,added test and adjust the param_grid

bac8e87

Muhammad-Rebaal and others added 3 commits July 21, 2026 13:50

Merge branch 'main' into optimization

8a8f97b

ENH: Tags are taken _BasePtForecasterV2, lightening tuner functions a…

4dcd61c

…dded code refactored and removed unwanted tests

Merge branch 'optimization' of https://github.com/Muhammad-Rebaal/pyt…

cdd7f85

…orch-forecasting into optimization

Muhammad-Rebaal added 3 commits August 26, 2026 19:51

HyperparameterTuner Feature Implemented

bed2247

Fix Code Quality

3c8d8d1

Fix Soft Depency issue

786ac9a

Muhammad-Rebaal marked this pull request as ready for review August 26, 2026 15:26

Muhammad-Rebaal and others added 5 commits August 27, 2026 00:29

Merge branch 'main' into optimization

52b28ab

HyperparameterTuner implemented without the need of datamodule_cfg an…

1470d12

…d trainer_cfg

Merge branch 'main' into optimization

c341a12

Added an edge case

ee7f848

Merge branch 'optimization' of https://github.com/Muhammad-Rebaal/pyt…

d02820a

…orch-forecasting into optimization

Muhammad-Rebaal and others added 3 commits September 8, 2026 17:50

datamodule_cls variable added into the HyperparameterTuner Class

7d03d39

Merge branch 'main' into optimization

c8f13d5

Merge branch 'main' into optimization

a2a2a7e

phoeenniixx reviewed Sep 10, 2026

View reviewed changes

phoeenniixx requested changes Sep 10, 2026

View reviewed changes

Muhammad-Rebaal and others added 5 commits September 11, 2026 18:21

Merge branch 'main' into optimization

2eaa29c

Merge branch 'main' into optimization

87f018e

Merge branch 'main' into optimization

f7662c3

cleaned up the test file and added the example

300ec2b

Merge branch 'optimization' of https://github.com/Muhammad-Rebaal/pyt…

f86629c

…orch-forecasting into optimization

Muhammad-Rebaal requested a review from phoeenniixx September 17, 2026 10:14

Merge branch 'main' into optimization

51e328c

phoeenniixx requested changes Sep 19, 2026

View reviewed changes

Muhammad-Rebaal added 4 commits September 25, 2026 21:07

Merge branch 'main' into optimization

51fe98f

Merge branch 'main' into optimization

78fd85c

Merge branch 'main' into optimization

f838efd

Merge branch 'main' into optimization

076338c

	super().__init__(
	loss=loss,
	logging_metrics=logging_metrics,
	optimizer=optimizer,
	optimizer_params=optimizer_params,
	lr_scheduler=lr_scheduler,
	lr_scheduler_params=lr_scheduler_params,
	)

Conversation

Muhammad-Rebaal commented Jul 7, 2026

Implemented Architecture

Uh oh!

codecov Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Jul 8, 2026

Uh oh!

fkiraly left a comment

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Jul 12, 2026

Uh oh!

phoeenniixx left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoeenniixx commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Muhammad-Rebaal commented Jul 14, 2026

Uh oh!

Muhammad-Rebaal commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Muhammad-Rebaal commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Muhammad-Rebaal commented Aug 31, 2026

Uh oh!

Muhammad-Rebaal commented Sep 10, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Sep 17, 2026

Uh oh!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoeenniixx Sep 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal Sep 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

codecov Bot commented Jul 7, 2026 •

edited

Loading

phoeenniixx left a comment •

edited

Loading

phoeenniixx commented Jul 13, 2026 •

edited

Loading

Muhammad-Rebaal commented Jul 15, 2026 •

edited

Loading

Muhammad-Rebaal commented Aug 26, 2026 •

edited

Loading

phoeenniixx Sep 28, 2026 •

edited

Loading

Muhammad-Rebaal Sep 28, 2026 •

edited

Loading