Customize some parameters for Study object (Optuna backend)

[wasf84]

Description

Hi.

It would be nice if it were possible to customize the Study object (Optuna backend) for all Auto models.

Sometimes the optimization process breaks with a 'RuntimeError' (or any other error). This happens because occasionally, depending on the parameters sampled for the model, a DataFrame full of null values is generated, and applying the validation metric causes an error. It would be possible to prevent the optimization process from stopping if the Study could be customized using the 'study.optimize(..., catch=(RuntimeError,))' attribute. The optimization process would simply ignore this error and proceed to the next one.

This could be achievied if we could pass some "study_kwargs" in the Auto* models creation, I think.

For example:

AutoDeepAR(
h=24,
config=config_deepar,
search_alg=optuna.samplers.TPESampler(seed=1989),
backend='optuna',
num_samples=5,
study_kwargs={'catch' : '(RuntimeError,)'}
)

I hope I made myself clear.

Thanks for your attention.

Use case

No response

[wasf84]

Hi!

I'm not sure if it's correct, but I've been trying to implement this feature by myself. I'd like to share with you guys what I've done. It's working, but as I said, I'm unsure if this is the correct way to do it.

from neuralforecast.common._base_auto import BaseAuto
from neuralforecast.models.tide import TiDE
from neuralforecast.losses.pytorch import MAE
from copy import deepcopy
import os
from ray.tune.search.basic_variant import BasicVariantGenerator

class CustomBaseAuto(BaseAuto):

    def __init__(
        self,
        cls_model,
        h,
        loss,
        valid_loss,
        config,
        search_alg=BasicVariantGenerator(random_state=1),
        num_samples=10,
        cpus=os.cpu_count(),
        gpus=torch.cuda.device_count(),
        refit_with_val=False,
        verbose=False,
        alias=None,
        backend="ray",
        callbacks=None,
        optimize_catch=None, # <<<<<< New parameter
    ):
        self.optimize_catch = optimize_catch # <<<<<< New parameter

        super().__init__(
            cls_model,
            h,
            loss,
            valid_loss,
            config,
            search_alg,
            num_samples,
            cpus,
            gpus,
            refit_with_val,
            verbose,
            alias,
            backend,
            callbacks,
        )

    def _optuna_tune_model(
        self,
        cls_model,
        dataset,
        val_size,
        test_size,
        verbose,
        num_samples,
        search_alg,
        config,
        distributed_config,
        optimize_catch, # <<<<<< New parameter
    ):
        import optuna

        def objective(trial):
            user_cfg = config(trial)
            cfg = deepcopy(user_cfg)
            model = self._fit_model(
                cls_model=cls_model,
                config=cfg,
                dataset=dataset,
                val_size=val_size,
                test_size=test_size,
                distributed_config=distributed_config,
            )
            trial.set_user_attr("ALL_PARAMS", user_cfg)
            metrics = model.metrics
            trial.set_user_attr(
                "METRICS",
                {
                    "loss": metrics["ptl/val_loss"],
                    "train_loss": metrics["train_loss"],
                },
            )
            return trial.user_attrs["METRICS"]["loss"]

        if isinstance(search_alg, optuna.samplers.BaseSampler):
            sampler = search_alg
        else:
            sampler = None

        study = optuna.create_study(sampler=sampler, direction="minimize")
        study.optimize(
            objective,
            n_trials=num_samples,
            show_progress_bar=verbose,
            callbacks=self.callbacks,
            catch=optimize_catch # <<<<<< New parameter
        )
        return study

    def fit(
        self,
        dataset,
        val_size=0,
        test_size=0,
        random_seed=None,
        distributed_config=None,
    ):

        # we need val_size > 0 to perform
        # hyperparameter selection.
        search_alg = deepcopy(self.search_alg)
        val_size = val_size if val_size > 0 else self.h
        if self.backend == "ray":
            if distributed_config is not None:
                raise ValueError(
                    "distributed training is not supported for the ray backend."
                )
            results = self._tune_model(
                cls_model=self.cls_model,
                dataset=dataset,
                val_size=val_size,
                test_size=test_size,
                cpus=self.cpus,
                gpus=self.gpus,
                verbose=self.verbose,
                num_samples=self.num_samples,
                search_alg=search_alg,
                config=self.config,
            )
            best_config = results.get_best_result().config
        else:
            results = self._optuna_tune_model(
                cls_model=self.cls_model,
                dataset=dataset,
                val_size=val_size,
                test_size=test_size,
                verbose=self.verbose,
                num_samples=self.num_samples,
                search_alg=search_alg,
                config=self.config,
                distributed_config=distributed_config,
                optimize_catch=self.optimize_catch # <<<<<< New parameter
            )
            best_config = results.best_trial.user_attrs["ALL_PARAMS"]
        self.model = self._fit_model(
            cls_model=self.cls_model,
            config=best_config,
            dataset=dataset,
            val_size=val_size * self.refit_with_val,
            test_size=test_size,
            distributed_config=distributed_config,
        )
        self.results = results

        # Added attributes for compatibility with NeuralForecast core
        self.futr_exog_list = self.model.futr_exog_list
        self.hist_exog_list = self.model.hist_exog_list
        self.stat_exog_list = self.model.stat_exog_list
        return self
# -----------------------
class CustomAutoTiDE(CustomBaseAuto):
    def __init__(
        self,
        h,
        loss=MAE(),
        valid_loss=None,
        config=None,
        search_alg=BasicVariantGenerator(random_state=1),
        num_samples=10,
        refit_with_val=False,
        cpus=os.cpu_count() // 2,
        gpus=torch.cuda.device_count(),
        verbose=False,
        alias=None,
        backend="ray",
        callbacks=None,
        optimize_catch=None, # <<<<<< New parameter
    ):

        super().__init__(
            cls_model=TiDE,
            h=h,
            loss=loss,
            valid_loss=valid_loss,
            config=config,
            search_alg=search_alg,
            num_samples=num_samples,
            refit_with_val=refit_with_val,
            cpus=cpus,
            gpus=gpus,
            verbose=verbose,
            alias=alias,
            backend=backend,
            callbacks=callbacks,
            optimize_catch=optimize_catch # <<<<<< New parameter
        )

Using this "new" model.

from neuralforecast.losses.pytorch import SMAPE

def config_customtide(trial):
    input_multiplier = trial.suggest_int("input_multiplier", 1, 12)

    config = {
        'random_seed' : SEED,
        'alias' : "TiDE-Opt",
        'logger' : False,
        'enable_progress_bar' : False,
        'enable_model_summary' : False,
        'scaler_type' : 'standard',
        'input_size' : input_multiplier * len(test),
        'num_encoder_layers' : trial.suggest_int("num_encoder_layers", 1, 6),
        'num_decoder_layers' : trial.suggest_int("num_decoder_layers", 1, 6),
        'learning_rate' : trial.suggest_float("learning_rate", 0.001, 0.2, step=0.001),
        'hidden_size' : trial.suggest_int("hidden_size", 16, 2048, step=2),
        'batch_size' : trial.suggest_int("batch_size", 4, 128, step=2),
    }

    return config

md = CustomAutoTiDE(
    h=len(test),
    loss=SMAPE(),
    config=config_customtide,
    search_alg=opt.samplers.TPESampler(seed=SEED),
    backend='optuna',
    num_samples=3,
    optimize_catch=(FloatingPointError, ValueError, RuntimeError), # <<<<<< New parameter
)

nf = NeuralForecast(models=[md], freq='D')
nf.fit(df=train[["ds", "unique_id", "y"]])

I haven't been able to confirm whether the new parameter is being passed correctly where I want, but with the way I structured the classes, I am no longer receiving the error messages I mentioned at the beginning of this post. Also, no issues occur during execution, so I believe I did it correctly. Does that sound right? Thank you.

[wasf84]

Hi guys.
Any thoughts? Am I doing this correctly?

Thank you.