merge_main

Nixtla · Nov 19, 2024 · c529ced · c529ced
2 parents ddc617f + 642ced4
commit c529ced
Show file tree

Hide file tree

Showing 73 changed files with 284 additions and 51 deletions.
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
@@ -33,7 +33,7 @@ jobs:
       - name: Build package
         run: python -m build
       - name: Publish package
-        uses: pypa/gh-action-pypi-publish@fb13cb306901256ace3dab689990e13a5550ffaa # v1.11.0
+        uses: pypa/gh-action-pypi-publish@15c56dba361d8335944d31a2ecd17d700fc7bcbc # v1.12.2
         with:
           user: __token__
           password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/nbs/common.base_model.ipynb b/nbs/common.base_model.ipynb
@@ -498,13 +498,25 @@
     "            datamodule_constructor = TimeSeriesDataModule\n",
     "        else:\n",
     "            datamodule_constructor = _DistributedTimeSeriesDataModule\n",
+    "        \n",
+    "        dataloader_kwargs = self.dataloader_kwargs if self.dataloader_kwargs is not None else {}\n",
+    "        \n",
+    "        if self.num_workers_loader != 0:  # value is not at its default\n",
+    "            warnings.warn(\n",
+    "                \"The `num_workers_loader` argument is deprecated and will be removed in a future version. \"\n",
+    "                \"Please provide num_workers through `dataloader_kwargs`, e.g. \"\n",
+    "                f\"`dataloader_kwargs={{'num_workers': {self.num_workers_loader}}}`\",\n",
+    "                category=FutureWarning,\n",
+    "            )\n",
+    "            dataloader_kwargs['num_workers'] = self.num_workers_loader\n",
+    "\n",
     "        datamodule = datamodule_constructor(\n",
     "            dataset=dataset, \n",
     "            batch_size=batch_size,\n",
     "            valid_batch_size=valid_batch_size,\n",
-    "            num_workers=self.num_workers_loader,\n",
     "            drop_last=self.drop_last_loader,\n",
     "            shuffle_train=shuffle_train,\n",
+    "            **dataloader_kwargs\n",
     "        )\n",
     "\n",
     "        if self.val_check_steps > self.max_steps:\n",

diff --git a/nbs/core.ipynb b/nbs/core.ipynb
@@ -1595,15 +1595,26 @@
     "        except FileNotFoundError:\n",
     "            raise Exception('No configuration found in directory.')\n",
     "\n",
+    "        # in 1.6.4, `local_scaler_type` / `scalers_` lived on the dataset.\n",
+    "        # in order to preserve backwards-compatibility, we check to see if these are found on the dataset\n",
+    "        # in case they cannot be found in `config_dict`\n",
+    "        default_scalar_type = getattr(dataset, \"local_scaler_type\", None)\n",
+    "        default_scalars_ = getattr(dataset, \"scalers_\", None)\n",
+    "\n",
     "        # Create NeuralForecast object\n",
     "        neuralforecast = NeuralForecast(\n",
     "            models=models,\n",
     "            freq=config_dict['freq'],\n",
-    "            local_scaler_type=config_dict['local_scaler_type'],\n",
+    "            local_scaler_type=config_dict.get(\"local_scaler_type\", default_scalar_type),\n",
     "        )\n",
     "\n",
-    "        for attr in ['id_col', 'time_col', 'target_col']:\n",
-    "            setattr(neuralforecast, attr, config_dict[attr])\n",
+    "        attr_to_default = {\n",
+    "            \"id_col\": \"unique_id\",\n",
+    "            \"time_col\": \"ds\",\n",
+    "            \"target_col\": \"y\"\n",
+    "        }\n",
+    "        for attr, default in attr_to_default.items():\n",
+    "            setattr(neuralforecast, attr, config_dict.get(attr, default))\n",
     "        # only restore attribute if available\n",
     "        for attr in ['prediction_intervals', '_cs_df']:\n",
     "            if attr in config_dict.keys():\n",
@@ -1624,7 +1635,7 @@
     "        # Fitted flag\n",
     "        neuralforecast._fitted = config_dict['_fitted']\n",
     "\n",
-    "        neuralforecast.scalers_ = config_dict['scalers_']\n",
+    "        neuralforecast.scalers_ = config_dict.get(\"scalers_\", default_scalars_)\n",
     "\n",
     "        return neuralforecast\n",
     "    \n",

diff --git a/nbs/docs/tutorials/18_adding_models.ipynb b/nbs/docs/tutorials/18_adding_models.ipynb
@@ -269,7 +269,6 @@
     "                 step_size: int = 1,\n",
     "                 scaler_type: str = 'identity',\n",
     "                 random_seed: int = 1,\n",
-    "                 num_workers_loader: int = 0,\n",
     "                 drop_last_loader: bool = False,\n",
     "                 **trainer_kwargs):\n",
     "    # Inherit BaseWindows class\n",
@@ -415,7 +414,13 @@
    ]
   }
  ],
- "metadata": {},
+ "metadata": {
+  "kernelspec": {
+   "display_name": "python3",
+   "language": "python",
+   "name": "python3"
+  }
+ },
  "nbformat": 4,
  "nbformat_minor": 2
 }
diff --git a/nbs/models.autoformer.ipynb b/nbs/models.autoformer.ipynb
@@ -466,6 +466,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>\n",
     "\n",
     "\t*References*<br>\n",
@@ -516,6 +517,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs=None,\n",
     "                 **trainer_kwargs):\n",
     "        super(Autoformer, self).__init__(h=h,\n",
     "                                       input_size=input_size,\n",
@@ -544,6 +546,7 @@
     "                                       optimizer_kwargs=optimizer_kwargs,\n",
     "                                       lr_scheduler=lr_scheduler,\n",
     "                                       lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "                                       dataloader_kwargs=dataloader_kwargs,\n",
     "                                       **trainer_kwargs)\n",
     "\n",
     "        # Architecture\n",

diff --git a/nbs/models.bitcn.ipynb b/nbs/models.bitcn.ipynb
@@ -185,6 +185,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>    \n",
     "\n",
     "    **References**<br>  \n",
@@ -228,6 +229,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs=None,\n",
     "                 **trainer_kwargs):\n",
     "        super(BiTCN, self).__init__(\n",
     "            h=h,\n",
@@ -257,6 +259,7 @@
     "            optimizer_kwargs=optimizer_kwargs,\n",
     "            lr_scheduler=lr_scheduler,\n",
     "            lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "            dataloader_kwargs=dataloader_kwargs,\n",
     "            **trainer_kwargs\n",
     "        )\n",
     "\n",

diff --git a/nbs/models.deepar.ipynb b/nbs/models.deepar.ipynb
@@ -186,6 +186,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>    \n",
     "\n",
     "    **References**<br>\n",
@@ -234,6 +235,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs = None,\n",
     "                 **trainer_kwargs):\n",
     "\n",
     "        if exclude_insample_y:\n",
@@ -267,6 +269,7 @@
     "                                    optimizer_kwargs=optimizer_kwargs,\n",
     "                                    lr_scheduler=lr_scheduler,\n",
     "                                    lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "                                    dataloader_kwargs=dataloader_kwargs,\n",
     "                                    **trainer_kwargs)\n",
     "\n",
     "        self.n_samples = trajectory_samples\n",

diff --git a/nbs/models.deepnpts.ipynb b/nbs/models.deepnpts.ipynb
@@ -127,6 +127,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>    \n",
     "\n",
     "    **References**<br>\n",
@@ -172,6 +173,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs = None,\n",
     "                 **trainer_kwargs):\n",
     "\n",
     "        if exclude_insample_y:\n",
@@ -211,6 +213,7 @@
     "                                    optimizer_kwargs=optimizer_kwargs,\n",
     "                                    lr_scheduler=lr_scheduler,\n",
     "                                    lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "                                    dataloader_kwargs=dataloader_kwargs,\n",
     "                                    **trainer_kwargs)\n",
     "\n",
     "        self.h = h\n",

diff --git a/nbs/models.dilated_rnn.ipynb b/nbs/models.dilated_rnn.ipynb
@@ -406,6 +406,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br> \n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>    \n",
     "    \"\"\"\n",
     "    # Class attributes\n",
@@ -450,6 +451,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs = None,\n",
     "                 **trainer_kwargs):\n",
     "        super(DilatedRNN, self).__init__(\n",
     "            h=h,\n",
@@ -479,6 +481,7 @@
     "            optimizer_kwargs=optimizer_kwargs,\n",
     "            lr_scheduler=lr_scheduler,\n",
     "            lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "            dataloader_kwargs=dataloader_kwargs,\n",
     "            **trainer_kwargs\n",
     "        )\n",
     "\n",

diff --git a/nbs/models.dlinear.ipynb b/nbs/models.dlinear.ipynb
@@ -169,6 +169,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>\n",
     "\n",
     "\t*References*<br>\n",
@@ -210,6 +211,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs=None,\n",
     "                 **trainer_kwargs):\n",
     "        super(DLinear, self).__init__(h=h,\n",
     "                                       input_size=input_size,\n",
@@ -238,6 +240,7 @@
     "                                       optimizer_kwargs=optimizer_kwargs,\n",
     "                                       lr_scheduler=lr_scheduler,\n",
     "                                       lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "                                       dataloader_kwargs=dataloader_kwargs,\n",
     "                                       **trainer_kwargs)\n",
     "                                                                \n",
     "        # Architecture\n",

diff --git a/nbs/models.fedformer.ipynb b/nbs/models.fedformer.ipynb
@@ -469,6 +469,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>\n",
     "\n",
     "    \"\"\"\n",
@@ -518,6 +519,7 @@
     "                 optimizer_kwargs=None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs = None,\n",
     "                 **trainer_kwargs):\n",
     "        super(FEDformer, self).__init__(h=h,\n",
     "                                       input_size=input_size,\n",
@@ -544,7 +546,8 @@
     "                                       optimizer=optimizer,\n",
     "                                       optimizer_kwargs=optimizer_kwargs,\n",
     "                                       lr_scheduler=lr_scheduler,\n",
-    "                                       lr_scheduler_kwargs=lr_scheduler_kwargs,                                    \n",
+    "                                       lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "                                       dataloader_kwargs=dataloader_kwargs,                                    \n",
     "                                       **trainer_kwargs)\n",
     "        # Architecture\n",
     "        self.label_len = int(np.ceil(input_size * decoder_input_size_multiplier))\n",

diff --git a/nbs/models.gru.ipynb b/nbs/models.gru.ipynb
@@ -132,6 +132,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>    \n",
     "    \"\"\"\n",
     "    # Class attributes\n",
@@ -179,6 +180,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs = None,\n",
     "                 **trainer_kwargs):\n",
     "        \n",
     "        self.RECURRENT = recurrent\n",
@@ -211,6 +213,7 @@
     "            optimizer_kwargs=optimizer_kwargs,\n",
     "            lr_scheduler=lr_scheduler,\n",
     "            lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "            dataloader_kwargs=dataloader_kwargs,\n",
     "            **trainer_kwargs\n",
     "        )\n",
     "\n",

diff --git a/nbs/models.informer.ipynb b/nbs/models.informer.ipynb
@@ -313,6 +313,7 @@
     "    `optimizer_kwargs`: dict, optional, list of parameters used by the user specified `optimizer`.<br>\n",
     "    `lr_scheduler`: Subclass of 'torch.optim.lr_scheduler.LRScheduler', optional, user specified lr_scheduler instead of the default choice (StepLR).<br>\n",
     "    `lr_scheduler_kwargs`: dict, optional, list of parameters used by the user specified `lr_scheduler`.<br>\n",
+    "    `dataloader_kwargs`: dict, optional, list of parameters passed into the PyTorch Lightning dataloader by the `TimeSeriesDataLoader`. <br>\n",
     "    `**trainer_kwargs`: int,  keyword trainer arguments inherited from [PyTorch Lighning's trainer](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.trainer.trainer.Trainer.html?highlight=trainer).<br>\n",
     "\n",
     "\t*References*<br>\n",
@@ -363,6 +364,7 @@
     "                 optimizer_kwargs = None,\n",
     "                 lr_scheduler = None,\n",
     "                 lr_scheduler_kwargs = None,\n",
+    "                 dataloader_kwargs = None,\n",
     "                 **trainer_kwargs):\n",
     "        super(Informer, self).__init__(h=h,\n",
     "                                       input_size=input_size,\n",
@@ -391,6 +393,7 @@
     "                                       optimizer_kwargs=optimizer_kwargs,\n",
     "                                       lr_scheduler=lr_scheduler,\n",
     "                                       lr_scheduler_kwargs=lr_scheduler_kwargs,\n",
+    "                                       dataloader_kwargs=dataloader_kwargs,\n",
     "                                       **trainer_kwargs)\n",
     "\n",
     "        # Architecture\n",