We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Default trainer kwargs in temporal_fusion_transformer/tuning.py cannot be overwritten
```python
[I 2021-01-27 09:42:03,226] A new study created in memory with name: no-name-8f54548f-a847-4450-9a1c-93acfdc5609e
[W 2021-01-27 09:42:03,229] Trial 0 failed because of the following error: TypeError("ABCMeta object got multiple values for keyword argument 'progress_bar_refresh_rate'",)
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/optuna/_optimize.py", line 198, in _run_trial
    value_or_values = func(trial)
  File "/usr/local/lib/python3.6/dist-packages/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py", line 145, in objective
    **trainer_kwargs,
TypeError: ABCMeta object got multiple values for keyword argument 'progress_bar_refresh_rate'
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_forecasting/models/temporal_fusion_transformer/tuning.py`
Content:
```
1 """
2 Hyperparameters can be efficiently tuned with `optuna <https://optuna.readthedocs.io/>`_.
3 """
4 import copy
5 import logging
6 import os
7 from typing import Any, Dict, Tuple, Union
8 
9 import numpy as np
10 import optuna
11 from optuna.integration import PyTorchLightningPruningCallback, TensorBoardCallback
12 import optuna.logging
13 import pytorch_lightning as pl
14 from pytorch_lightning import Callback
15 from pytorch_lightning.callbacks import LearningRateMonitor
16 from pytorch_lightning.loggers import TensorBoardLogger
17 import statsmodels.api as sm
18 import torch
19 from torch.utils.data import DataLoader
20 
21 from pytorch_forecasting import TemporalFusionTransformer
22 from pytorch_forecasting.data import TimeSeriesDataSet
23 from pytorch_forecasting.metrics import QuantileLoss
24 
25 optuna_logger = logging.getLogger("optuna")
26 
27 
28 class MetricsCallback(Callback):
29     """PyTorch Lightning metric callback."""
30 
31     def __init__(self):
32         super().__init__()
33         self.metrics = []
34 
35     def on_validation_end(self, trainer, pl_module):
36         self.metrics.append(trainer.callback_metrics)
37 
38 
39 def optimize_hyperparameters(
40     train_dataloader: DataLoader,
41     val_dataloader: DataLoader,
42     model_path: str,
43     max_epochs: int = 20,
44     n_trials: int = 100,
45     timeout: float = 3600 * 8.0,  # 8 hours
46     gradient_clip_val_range: Tuple[float, float] = (0.01, 100.0),
47     hidden_size_range: Tuple[int, int] = (16, 265),
48     hidden_continuous_size_range: Tuple[int, int] = (8, 64),
49     attention_head_size_range: Tuple[int, int] = (1, 4),
50     dropout_range: Tuple[float, float] = (0.1, 0.3),
51     learning_rate_range: Tuple[float, float] = (1e-5, 1.0),
52     use_learning_rate_finder: bool = True,
53     trainer_kwargs: Dict[str, Any] = {},
54     log_dir: str = "lightning_logs",
55     study: optuna.Study = None,
56     verbose: Union[int, bool] = None,
57     **kwargs,
58 ) -> optuna.Study:
59     """
60     Optimize Temporal Fusion Transformer hyperparameters.
61 
62     Run hyperparameter optimization. Learning rate for is determined with
63     the PyTorch Lightning learning rate finder.
64 
65     Args:
66         train_dataloader (DataLoader): dataloader for training model
67         val_dataloader (DataLoader): dataloader for validating model
68         model_path (str): folder to which model checkpoints are saved
69         max_epochs (int, optional): Maximum number of epochs to run training. Defaults to 20.
70         n_trials (int, optional): Number of hyperparameter trials to run. Defaults to 100.
71         timeout (float, optional): Time in seconds after which training is stopped regardless of number of epochs
72             or validation metric. Defaults to 3600*8.0.
73         hidden_size_range (Tuple[int, int], optional): Minimum and maximum of ``hidden_size`` hyperparameter. Defaults
74             to (16, 265).
75         hidden_continuous_size_range (Tuple[int, int], optional):  Minimum and maximum of ``hidden_continuous_size``
76             hyperparameter. Defaults to (8, 64).
77         attention_head_size_range (Tuple[int, int], optional):  Minimum and maximum of ``attention_head_size``
78             hyperparameter. Defaults to (1, 4).
79         dropout_range (Tuple[float, float], optional):  Minimum and maximum of ``dropout`` hyperparameter. Defaults to
80             (0.1, 0.3).
81         learning_rate_range (Tuple[float, float], optional): Learning rate range. Defaults to (1e-5, 1.0).
82         use_learning_rate_finder (bool): If to use learning rate finder or optimize as part of hyperparameters.
83             Defaults to True.
84         trainer_kwargs (Dict[str, Any], optional): Additional arguments to the
85             `PyTorch Lightning trainer <https://pytorch-lightning.readthedocs.io/en/latest/trainer.html>`_ such
86             as ``limit_train_batches``. Defaults to {}.
87         log_dir (str, optional): Folder into which to log results for tensorboard. Defaults to "lightning_logs".
88         study (optuna.Study, optional): study to resume. Will create new study by default.
89         verbose (Union[int, bool]): level of verbosity.
90             * None: no change in verbosity level (equivalent to verbose=1 by optuna-set default).
91             * 0 or False: log only warnings.
92             * 1 or True: log pruning events.
93             * 2: optuna logging level at debug level.
94             Defaults to None.
95 
96         **kwargs: Additional arguments for the :py:class:`~TemporalFusionTransformer`.
97 
98     Returns:
99         optuna.Study: optuna study results
100     """
101     assert isinstance(train_dataloader.dataset, TimeSeriesDataSet) and isinstance(
102         val_dataloader.dataset, TimeSeriesDataSet
103     ), "dataloaders must be built from timeseriesdataset"
104 
105     logging_level = {
106         None: optuna.logging.get_verbosity(),
107         0: optuna.logging.WARNING,
108         1: optuna.logging.INFO,
109         2: optuna.logging.DEBUG,
110     }
111     optuna_verbose = logging_level[verbose]
112     optuna.logging.set_verbosity(optuna_verbose)
113 
114     loss = kwargs.get(
115         "loss", QuantileLoss()
116     )  # need a deepcopy of loss as it will otherwise propagate from one trial to the next
117 
118     # create objective function
119     def objective(trial: optuna.Trial) -> float:
120         # Filenames for each trial must be made unique in order to access each checkpoint.
121         checkpoint_callback = pl.callbacks.ModelCheckpoint(
122             dirpath=os.path.join(model_path, "trial_{}".format(trial.number)), filename="{epoch}", monitor="val_loss"
123         )
124 
125         # The default logger in PyTorch Lightning writes to event files to be consumed by
126         # TensorBoard. We don't use any logger here as it requires us to implement several abstract
127         # methods. Instead we setup a simple callback, that saves metrics from each validation step.
128         metrics_callback = MetricsCallback()
129         learning_rate_callback = LearningRateMonitor()
130         logger = TensorBoardLogger(log_dir, name="optuna", version=trial.number)
131         gradient_clip_val = trial.suggest_loguniform("gradient_clip_val", *gradient_clip_val_range)
132         trainer_kwargs.setdefault("gpus", [0] if torch.cuda.is_available() else None)
133         trainer = pl.Trainer(
134             max_epochs=max_epochs,
135             gradient_clip_val=gradient_clip_val,
136             callbacks=[
137                 metrics_callback,
138                 learning_rate_callback,
139                 checkpoint_callback,
140                 PyTorchLightningPruningCallback(trial, monitor="val_loss"),
141             ],
142             logger=logger,
143             progress_bar_refresh_rate=[0, 1][optuna_verbose < optuna.logging.INFO],
144             weights_summary=[None, "top"][optuna_verbose < optuna.logging.INFO],
145             **trainer_kwargs,
146         )
147 
148         # create model
149         hidden_size = trial.suggest_int("hidden_size", *hidden_size_range, log=True)
150         kwargs["loss"] = copy.deepcopy(loss)
151         model = TemporalFusionTransformer.from_dataset(
152             train_dataloader.dataset,
153             dropout=trial.suggest_uniform("dropout", *dropout_range),
154             hidden_size=hidden_size,
155             hidden_continuous_size=trial.suggest_int(
156                 "hidden_continuous_size",
157                 hidden_continuous_size_range[0],
158                 min(hidden_continuous_size_range[1], hidden_size),
159                 log=True,
160             ),
161             attention_head_size=trial.suggest_int("attention_head_size", *attention_head_size_range),
162             log_interval=-1,
163             **kwargs,
164         )
165         # find good learning rate
166         if use_learning_rate_finder:
167             lr_trainer = pl.Trainer(
168                 gradient_clip_val=gradient_clip_val,
169                 gpus=[0] if torch.cuda.is_available() else None,
170                 logger=False,
171                 progress_bar_refresh_rate=0,
172                 weights_summary=None,
173             )
174             res = lr_trainer.tuner.lr_find(
175                 model,
176                 train_dataloader=train_dataloader,
177                 val_dataloaders=val_dataloader,
178                 early_stop_threshold=10000,
179                 min_lr=learning_rate_range[0],
180                 num_training=100,
181                 max_lr=learning_rate_range[1],
182             )
183 
184             loss_finite = np.isfinite(res.results["loss"])
185             if loss_finite.sum() > 3:  # at least 3 valid values required for learning rate finder
186                 lr_smoothed, loss_smoothed = sm.nonparametric.lowess(
187                     np.asarray(res.results["loss"])[loss_finite],
188                     np.asarray(res.results["lr"])[loss_finite],
189                     frac=1.0 / 10.0,
190                 )[min(loss_finite.sum() - 3, 10) : -1].T
191                 optimal_idx = np.gradient(loss_smoothed).argmin()
192                 optimal_lr = lr_smoothed[optimal_idx]
193             else:
194                 optimal_idx = np.asarray(res.results["loss"]).argmin()
195                 optimal_lr = res.results["lr"][optimal_idx]
196             optuna_logger.info(f"Using learning rate of {optimal_lr:.3g}")
197             # add learning rate artificially
198             model.hparams.learning_rate = trial.suggest_uniform("learning_rate", optimal_lr, optimal_lr)
199         else:
200             model.hparams.learning_rate = trial.suggest_loguniform("learning_rate", *learning_rate_range)
201 
202         # fit
203         trainer.fit(model, train_dataloader=train_dataloader, val_dataloaders=val_dataloader)
204 
205         # report result
206         return metrics_callback.metrics[-1]["val_loss"].item()
207 
208     # setup optuna and run
209     pruner = optuna.pruners.SuccessiveHalvingPruner()
210     if study is None:
211         study = optuna.create_study(direction="minimize", pruner=pruner)
212     study.optimize(objective, n_trials=n_trials, timeout=timeout)
213     return study
214
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py b/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py
--- a/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py
+++ b/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py
@@ -129,8 +129,8 @@
         learning_rate_callback = LearningRateMonitor()
         logger = TensorBoardLogger(log_dir, name="optuna", version=trial.number)
         gradient_clip_val = trial.suggest_loguniform("gradient_clip_val", *gradient_clip_val_range)
-        trainer_kwargs.setdefault("gpus", [0] if torch.cuda.is_available() else None)
-        trainer = pl.Trainer(
+        default_trainer_kwargs = dict(
+            gpus=[0] if torch.cuda.is_available() else None,
             max_epochs=max_epochs,
             gradient_clip_val=gradient_clip_val,
             callbacks=[
@@ -142,7 +142,10 @@
             logger=logger,
             progress_bar_refresh_rate=[0, 1][optuna_verbose < optuna.logging.INFO],
             weights_summary=[None, "top"][optuna_verbose < optuna.logging.INFO],
-            **trainer_kwargs,
+        )
+        default_trainer_kwargs.update(trainer_kwargs)
+        trainer = pl.Trainer(
+            **default_trainer_kwargs,
         )
 
         # create model

{"golden_diff": "diff --git a/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py b/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py\n--- a/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py\n+++ b/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py\n@@ -129,8 +129,8 @@\n         learning_rate_callback = LearningRateMonitor()\n         logger = TensorBoardLogger(log_dir, name=\"optuna\", version=trial.number)\n         gradient_clip_val = trial.suggest_loguniform(\"gradient_clip_val\", *gradient_clip_val_range)\n-        trainer_kwargs.setdefault(\"gpus\", [0] if torch.cuda.is_available() else None)\n-        trainer = pl.Trainer(\n+        default_trainer_kwargs = dict(\n+            gpus=[0] if torch.cuda.is_available() else None,\n             max_epochs=max_epochs,\n             gradient_clip_val=gradient_clip_val,\n             callbacks=[\n@@ -142,7 +142,10 @@\n             logger=logger,\n             progress_bar_refresh_rate=[0, 1][optuna_verbose < optuna.logging.INFO],\n             weights_summary=[None, \"top\"][optuna_verbose < optuna.logging.INFO],\n-            **trainer_kwargs,\n+        )\n+        default_trainer_kwargs.update(trainer_kwargs)\n+        trainer = pl.Trainer(\n+            **default_trainer_kwargs,\n         )\n \n         # create model\n", "issue": "Default trainer kwargs in temporal_fusion_transformer/tuning.py cannot be overwritten\n```python\r\n[I 2021-01-27 09:42:03,226] A new study created in memory with name: no-name-8f54548f-a847-4450-9a1c-93acfdc5609e\r\n[W 2021-01-27 09:42:03,229] Trial 0 failed because of the following error: TypeError(\"ABCMeta object got multiple values for keyword argument 'progress_bar_refresh_rate'\",)\r\nTraceback (most recent call last):\r\n  File \"/usr/local/lib/python3.6/dist-packages/optuna/_optimize.py\", line 198, in _run_trial\r\n    value_or_values = func(trial)\r\n  File \"/usr/local/lib/python3.6/dist-packages/pytorch_forecasting/models/temporal_fusion_transformer/tuning.py\", line 145, in objective\r\n    **trainer_kwargs,\r\nTypeError: ABCMeta object got multiple values for keyword argument 'progress_bar_refresh_rate'\r\n```\n", "before_files": [{"content": "\"\"\"\nHyperparameters can be efficiently tuned with `optuna <https://optuna.readthedocs.io/>`_.\n\"\"\"\nimport copy\nimport logging\nimport os\nfrom typing import Any, Dict, Tuple, Union\n\nimport numpy as np\nimport optuna\nfrom optuna.integration import PyTorchLightningPruningCallback, TensorBoardCallback\nimport optuna.logging\nimport pytorch_lightning as pl\nfrom pytorch_lightning import Callback\nfrom pytorch_lightning.callbacks import LearningRateMonitor\nfrom pytorch_lightning.loggers import TensorBoardLogger\nimport statsmodels.api as sm\nimport torch\nfrom torch.utils.data import DataLoader\n\nfrom pytorch_forecasting import TemporalFusionTransformer\nfrom pytorch_forecasting.data import TimeSeriesDataSet\nfrom pytorch_forecasting.metrics import QuantileLoss\n\noptuna_logger = logging.getLogger(\"optuna\")\n\n\nclass MetricsCallback(Callback):\n    \"\"\"PyTorch Lightning metric callback.\"\"\"\n\n    def __init__(self):\n        super().__init__()\n        self.metrics = []\n\n    def on_validation_end(self, trainer, pl_module):\n        self.metrics.append(trainer.callback_metrics)\n\n\ndef optimize_hyperparameters(\n    train_dataloader: DataLoader,\n    val_dataloader: DataLoader,\n    model_path: str,\n    max_epochs: int = 20,\n    n_trials: int = 100,\n    timeout: float = 3600 * 8.0,  # 8 hours\n    gradient_clip_val_range: Tuple[float, float] = (0.01, 100.0),\n    hidden_size_range: Tuple[int, int] = (16, 265),\n    hidden_continuous_size_range: Tuple[int, int] = (8, 64),\n    attention_head_size_range: Tuple[int, int] = (1, 4),\n    dropout_range: Tuple[float, float] = (0.1, 0.3),\n    learning_rate_range: Tuple[float, float] = (1e-5, 1.0),\n    use_learning_rate_finder: bool = True,\n    trainer_kwargs: Dict[str, Any] = {},\n    log_dir: str = \"lightning_logs\",\n    study: optuna.Study = None,\n    verbose: Union[int, bool] = None,\n    **kwargs,\n) -> optuna.Study:\n    \"\"\"\n    Optimize Temporal Fusion Transformer hyperparameters.\n\n    Run hyperparameter optimization. Learning rate for is determined with\n    the PyTorch Lightning learning rate finder.\n\n    Args:\n        train_dataloader (DataLoader): dataloader for training model\n        val_dataloader (DataLoader): dataloader for validating model\n        model_path (str): folder to which model checkpoints are saved\n        max_epochs (int, optional): Maximum number of epochs to run training. Defaults to 20.\n        n_trials (int, optional): Number of hyperparameter trials to run. Defaults to 100.\n        timeout (float, optional): Time in seconds after which training is stopped regardless of number of epochs\n            or validation metric. Defaults to 3600*8.0.\n        hidden_size_range (Tuple[int, int], optional): Minimum and maximum of ``hidden_size`` hyperparameter. Defaults\n            to (16, 265).\n        hidden_continuous_size_range (Tuple[int, int], optional):  Minimum and maximum of ``hidden_continuous_size``\n            hyperparameter. Defaults to (8, 64).\n        attention_head_size_range (Tuple[int, int], optional):  Minimum and maximum of ``attention_head_size``\n            hyperparameter. Defaults to (1, 4).\n        dropout_range (Tuple[float, float], optional):  Minimum and maximum of ``dropout`` hyperparameter. Defaults to\n            (0.1, 0.3).\n        learning_rate_range (Tuple[float, float], optional): Learning rate range. Defaults to (1e-5, 1.0).\n        use_learning_rate_finder (bool): If to use learning rate finder or optimize as part of hyperparameters.\n            Defaults to True.\n        trainer_kwargs (Dict[str, Any], optional): Additional arguments to the\n            `PyTorch Lightning trainer <https://pytorch-lightning.readthedocs.io/en/latest/trainer.html>`_ such\n            as ``limit_train_batches``. Defaults to {}.\n        log_dir (str, optional): Folder into which to log results for tensorboard. Defaults to \"lightning_logs\".\n        study (optuna.Study, optional): study to resume. Will create new study by default.\n        verbose (Union[int, bool]): level of verbosity.\n            * None: no change in verbosity level (equivalent to verbose=1 by optuna-set default).\n            * 0 or False: log only warnings.\n            * 1 or True: log pruning events.\n            * 2: optuna logging level at debug level.\n            Defaults to None.\n\n        **kwargs: Additional arguments for the :py:class:`~TemporalFusionTransformer`.\n\n    Returns:\n        optuna.Study: optuna study results\n    \"\"\"\n    assert isinstance(train_dataloader.dataset, TimeSeriesDataSet) and isinstance(\n        val_dataloader.dataset, TimeSeriesDataSet\n    ), \"dataloaders must be built from timeseriesdataset\"\n\n    logging_level = {\n        None: optuna.logging.get_verbosity(),\n        0: optuna.logging.WARNING,\n        1: optuna.logging.INFO,\n        2: optuna.logging.DEBUG,\n    }\n    optuna_verbose = logging_level[verbose]\n    optuna.logging.set_verbosity(optuna_verbose)\n\n    loss = kwargs.get(\n        \"loss\", QuantileLoss()\n    )  # need a deepcopy of loss as it will otherwise propagate from one trial to the next\n\n    # create objective function\n    def objective(trial: optuna.Trial) -> float:\n        # Filenames for each trial must be made unique in order to access each checkpoint.\n        checkpoint_callback = pl.callbacks.ModelCheckpoint(\n            dirpath=os.path.join(model_path, \"trial_{}\".format(trial.number)), filename=\"{epoch}\", monitor=\"val_loss\"\n        )\n\n        # The default logger in PyTorch Lightning writes to event files to be consumed by\n        # TensorBoard. We don't use any logger here as it requires us to implement several abstract\n        # methods. Instead we setup a simple callback, that saves metrics from each validation step.\n        metrics_callback = MetricsCallback()\n        learning_rate_callback = LearningRateMonitor()\n        logger = TensorBoardLogger(log_dir, name=\"optuna\", version=trial.number)\n        gradient_clip_val = trial.suggest_loguniform(\"gradient_clip_val\", *gradient_clip_val_range)\n        trainer_kwargs.setdefault(\"gpus\", [0] if torch.cuda.is_available() else None)\n        trainer = pl.Trainer(\n            max_epochs=max_epochs,\n            gradient_clip_val=gradient_clip_val,\n            callbacks=[\n                metrics_callback,\n                learning_rate_callback,\n                checkpoint_callback,\n                PyTorchLightningPruningCallback(trial, monitor=\"val_loss\"),\n            ],\n            logger=logger,\n            progress_bar_refresh_rate=[0, 1][optuna_verbose < optuna.logging.INFO],\n            weights_summary=[None, \"top\"][optuna_verbose < optuna.logging.INFO],\n            **trainer_kwargs,\n        )\n\n        # create model\n        hidden_size = trial.suggest_int(\"hidden_size\", *hidden_size_range, log=True)\n        kwargs[\"loss\"] = copy.deepcopy(loss)\n        model = TemporalFusionTransformer.from_dataset(\n            train_dataloader.dataset,\n            dropout=trial.suggest_uniform(\"dropout\", *dropout_range),\n            hidden_size=hidden_size,\n            hidden_continuous_size=trial.suggest_int(\n                \"hidden_continuous_size\",\n                hidden_continuous_size_range[0],\n                min(hidden_continuous_size_range[1], hidden_size),\n                log=True,\n            ),\n            attention_head_size=trial.suggest_int(\"attention_head_size\", *attention_head_size_range),\n            log_interval=-1,\n            **kwargs,\n        )\n        # find good learning rate\n        if use_learning_rate_finder:\n            lr_trainer = pl.Trainer(\n                gradient_clip_val=gradient_clip_val,\n                gpus=[0] if torch.cuda.is_available() else None,\n                logger=False,\n                progress_bar_refresh_rate=0,\n                weights_summary=None,\n            )\n            res = lr_trainer.tuner.lr_find(\n                model,\n                train_dataloader=train_dataloader,\n                val_dataloaders=val_dataloader,\n                early_stop_threshold=10000,\n                min_lr=learning_rate_range[0],\n                num_training=100,\n                max_lr=learning_rate_range[1],\n            )\n\n            loss_finite = np.isfinite(res.results[\"loss\"])\n            if loss_finite.sum() > 3:  # at least 3 valid values required for learning rate finder\n                lr_smoothed, loss_smoothed = sm.nonparametric.lowess(\n                    np.asarray(res.results[\"loss\"])[loss_finite],\n                    np.asarray(res.results[\"lr\"])[loss_finite],\n                    frac=1.0 / 10.0,\n                )[min(loss_finite.sum() - 3, 10) : -1].T\n                optimal_idx = np.gradient(loss_smoothed).argmin()\n                optimal_lr = lr_smoothed[optimal_idx]\n            else:\n                optimal_idx = np.asarray(res.results[\"loss\"]).argmin()\n                optimal_lr = res.results[\"lr\"][optimal_idx]\n            optuna_logger.info(f\"Using learning rate of {optimal_lr:.3g}\")\n            # add learning rate artificially\n            model.hparams.learning_rate = trial.suggest_uniform(\"learning_rate\", optimal_lr, optimal_lr)\n        else:\n            model.hparams.learning_rate = trial.suggest_loguniform(\"learning_rate\", *learning_rate_range)\n\n        # fit\n        trainer.fit(model, train_dataloader=train_dataloader, val_dataloaders=val_dataloader)\n\n        # report result\n        return metrics_callback.metrics[-1][\"val_loss\"].item()\n\n    # setup optuna and run\n    pruner = optuna.pruners.SuccessiveHalvingPruner()\n    if study is None:\n        study = optuna.create_study(direction=\"minimize\", pruner=pruner)\n    study.optimize(objective, n_trials=n_trials, timeout=timeout)\n    return study\n", "path": "pytorch_forecasting/models/temporal_fusion_transformer/tuning.py"}], "after_files": [{"content": "\"\"\"\nHyperparameters can be efficiently tuned with `optuna <https://optuna.readthedocs.io/>`_.\n\"\"\"\nimport copy\nimport logging\nimport os\nfrom typing import Any, Dict, Tuple, Union\n\nimport numpy as np\nimport optuna\nfrom optuna.integration import PyTorchLightningPruningCallback, TensorBoardCallback\nimport optuna.logging\nimport pytorch_lightning as pl\nfrom pytorch_lightning import Callback\nfrom pytorch_lightning.callbacks import LearningRateMonitor\nfrom pytorch_lightning.loggers import TensorBoardLogger\nimport statsmodels.api as sm\nimport torch\nfrom torch.utils.data import DataLoader\n\nfrom pytorch_forecasting import TemporalFusionTransformer\nfrom pytorch_forecasting.data import TimeSeriesDataSet\nfrom pytorch_forecasting.metrics import QuantileLoss\n\noptuna_logger = logging.getLogger(\"optuna\")\n\n\nclass MetricsCallback(Callback):\n    \"\"\"PyTorch Lightning metric callback.\"\"\"\n\n    def __init__(self):\n        super().__init__()\n        self.metrics = []\n\n    def on_validation_end(self, trainer, pl_module):\n        self.metrics.append(trainer.callback_metrics)\n\n\ndef optimize_hyperparameters(\n    train_dataloader: DataLoader,\n    val_dataloader: DataLoader,\n    model_path: str,\n    max_epochs: int = 20,\n    n_trials: int = 100,\n    timeout: float = 3600 * 8.0,  # 8 hours\n    gradient_clip_val_range: Tuple[float, float] = (0.01, 100.0),\n    hidden_size_range: Tuple[int, int] = (16, 265),\n    hidden_continuous_size_range: Tuple[int, int] = (8, 64),\n    attention_head_size_range: Tuple[int, int] = (1, 4),\n    dropout_range: Tuple[float, float] = (0.1, 0.3),\n    learning_rate_range: Tuple[float, float] = (1e-5, 1.0),\n    use_learning_rate_finder: bool = True,\n    trainer_kwargs: Dict[str, Any] = {},\n    log_dir: str = \"lightning_logs\",\n    study: optuna.Study = None,\n    verbose: Union[int, bool] = None,\n    **kwargs,\n) -> optuna.Study:\n    \"\"\"\n    Optimize Temporal Fusion Transformer hyperparameters.\n\n    Run hyperparameter optimization. Learning rate for is determined with\n    the PyTorch Lightning learning rate finder.\n\n    Args:\n        train_dataloader (DataLoader): dataloader for training model\n        val_dataloader (DataLoader): dataloader for validating model\n        model_path (str): folder to which model checkpoints are saved\n        max_epochs (int, optional): Maximum number of epochs to run training. Defaults to 20.\n        n_trials (int, optional): Number of hyperparameter trials to run. Defaults to 100.\n        timeout (float, optional): Time in seconds after which training is stopped regardless of number of epochs\n            or validation metric. Defaults to 3600*8.0.\n        hidden_size_range (Tuple[int, int], optional): Minimum and maximum of ``hidden_size`` hyperparameter. Defaults\n            to (16, 265).\n        hidden_continuous_size_range (Tuple[int, int], optional):  Minimum and maximum of ``hidden_continuous_size``\n            hyperparameter. Defaults to (8, 64).\n        attention_head_size_range (Tuple[int, int], optional):  Minimum and maximum of ``attention_head_size``\n            hyperparameter. Defaults to (1, 4).\n        dropout_range (Tuple[float, float], optional):  Minimum and maximum of ``dropout`` hyperparameter. Defaults to\n            (0.1, 0.3).\n        learning_rate_range (Tuple[float, float], optional): Learning rate range. Defaults to (1e-5, 1.0).\n        use_learning_rate_finder (bool): If to use learning rate finder or optimize as part of hyperparameters.\n            Defaults to True.\n        trainer_kwargs (Dict[str, Any], optional): Additional arguments to the\n            `PyTorch Lightning trainer <https://pytorch-lightning.readthedocs.io/en/latest/trainer.html>`_ such\n            as ``limit_train_batches``. Defaults to {}.\n        log_dir (str, optional): Folder into which to log results for tensorboard. Defaults to \"lightning_logs\".\n        study (optuna.Study, optional): study to resume. Will create new study by default.\n        verbose (Union[int, bool]): level of verbosity.\n            * None: no change in verbosity level (equivalent to verbose=1 by optuna-set default).\n            * 0 or False: log only warnings.\n            * 1 or True: log pruning events.\n            * 2: optuna logging level at debug level.\n            Defaults to None.\n\n        **kwargs: Additional arguments for the :py:class:`~TemporalFusionTransformer`.\n\n    Returns:\n        optuna.Study: optuna study results\n    \"\"\"\n    assert isinstance(train_dataloader.dataset, TimeSeriesDataSet) and isinstance(\n        val_dataloader.dataset, TimeSeriesDataSet\n    ), \"dataloaders must be built from timeseriesdataset\"\n\n    logging_level = {\n        None: optuna.logging.get_verbosity(),\n        0: optuna.logging.WARNING,\n        1: optuna.logging.INFO,\n        2: optuna.logging.DEBUG,\n    }\n    optuna_verbose = logging_level[verbose]\n    optuna.logging.set_verbosity(optuna_verbose)\n\n    loss = kwargs.get(\n        \"loss\", QuantileLoss()\n    )  # need a deepcopy of loss as it will otherwise propagate from one trial to the next\n\n    # create objective function\n    def objective(trial: optuna.Trial) -> float:\n        # Filenames for each trial must be made unique in order to access each checkpoint.\n        checkpoint_callback = pl.callbacks.ModelCheckpoint(\n            dirpath=os.path.join(model_path, \"trial_{}\".format(trial.number)), filename=\"{epoch}\", monitor=\"val_loss\"\n        )\n\n        # The default logger in PyTorch Lightning writes to event files to be consumed by\n        # TensorBoard. We don't use any logger here as it requires us to implement several abstract\n        # methods. Instead we setup a simple callback, that saves metrics from each validation step.\n        metrics_callback = MetricsCallback()\n        learning_rate_callback = LearningRateMonitor()\n        logger = TensorBoardLogger(log_dir, name=\"optuna\", version=trial.number)\n        gradient_clip_val = trial.suggest_loguniform(\"gradient_clip_val\", *gradient_clip_val_range)\n        default_trainer_kwargs = dict(\n            gpus=[0] if torch.cuda.is_available() else None,\n            max_epochs=max_epochs,\n            gradient_clip_val=gradient_clip_val,\n            callbacks=[\n                metrics_callback,\n                learning_rate_callback,\n                checkpoint_callback,\n                PyTorchLightningPruningCallback(trial, monitor=\"val_loss\"),\n            ],\n            logger=logger,\n            progress_bar_refresh_rate=[0, 1][optuna_verbose < optuna.logging.INFO],\n            weights_summary=[None, \"top\"][optuna_verbose < optuna.logging.INFO],\n        )\n        default_trainer_kwargs.update(trainer_kwargs)\n        trainer = pl.Trainer(\n            **default_trainer_kwargs,\n        )\n\n        # create model\n        hidden_size = trial.suggest_int(\"hidden_size\", *hidden_size_range, log=True)\n        kwargs[\"loss\"] = copy.deepcopy(loss)\n        model = TemporalFusionTransformer.from_dataset(\n            train_dataloader.dataset,\n            dropout=trial.suggest_uniform(\"dropout\", *dropout_range),\n            hidden_size=hidden_size,\n            hidden_continuous_size=trial.suggest_int(\n                \"hidden_continuous_size\",\n                hidden_continuous_size_range[0],\n                min(hidden_continuous_size_range[1], hidden_size),\n                log=True,\n            ),\n            attention_head_size=trial.suggest_int(\"attention_head_size\", *attention_head_size_range),\n            log_interval=-1,\n            **kwargs,\n        )\n        # find good learning rate\n        if use_learning_rate_finder:\n            lr_trainer = pl.Trainer(\n                gradient_clip_val=gradient_clip_val,\n                gpus=[0] if torch.cuda.is_available() else None,\n                logger=False,\n                progress_bar_refresh_rate=0,\n                weights_summary=None,\n            )\n            res = lr_trainer.tuner.lr_find(\n                model,\n                train_dataloader=train_dataloader,\n                val_dataloaders=val_dataloader,\n                early_stop_threshold=10000,\n                min_lr=learning_rate_range[0],\n                num_training=100,\n                max_lr=learning_rate_range[1],\n            )\n\n            loss_finite = np.isfinite(res.results[\"loss\"])\n            if loss_finite.sum() > 3:  # at least 3 valid values required for learning rate finder\n                lr_smoothed, loss_smoothed = sm.nonparametric.lowess(\n                    np.asarray(res.results[\"loss\"])[loss_finite],\n                    np.asarray(res.results[\"lr\"])[loss_finite],\n                    frac=1.0 / 10.0,\n                )[min(loss_finite.sum() - 3, 10) : -1].T\n                optimal_idx = np.gradient(loss_smoothed).argmin()\n                optimal_lr = lr_smoothed[optimal_idx]\n            else:\n                optimal_idx = np.asarray(res.results[\"loss\"]).argmin()\n                optimal_lr = res.results[\"lr\"][optimal_idx]\n            optuna_logger.info(f\"Using learning rate of {optimal_lr:.3g}\")\n            # add learning rate artificially\n            model.hparams.learning_rate = trial.suggest_uniform(\"learning_rate\", optimal_lr, optimal_lr)\n        else:\n            model.hparams.learning_rate = trial.suggest_loguniform(\"learning_rate\", *learning_rate_range)\n\n        # fit\n        trainer.fit(model, train_dataloader=train_dataloader, val_dataloaders=val_dataloader)\n\n        # report result\n        return metrics_callback.metrics[-1][\"val_loss\"].item()\n\n    # setup optuna and run\n    pruner = optuna.pruners.SuccessiveHalvingPruner()\n    if study is None:\n        study = optuna.create_study(direction=\"minimize\", pruner=pruner)\n    study.optimize(objective, n_trials=n_trials, timeout=timeout)\n    return study\n", "path": "pytorch_forecasting/models/temporal_fusion_transformer/tuning.py"}]}