Merge branch 'main' into feat/modify-config-optimizers

Nixtla · Dec 11, 2024 · c0a7eb8 · c0a7eb8
2 parents 3539eb3 + 4895284
commit c0a7eb8
Show file tree

Hide file tree

Showing 108 changed files with 4,836 additions and 1,218 deletions.
diff --git a/.all-contributorsrc b/.all-contributorsrc
@@ -6,10 +6,10 @@
   "commit": false,
   "contributors": [
     {
-      "login": "FedericoGarza",
-      "name": "fede",
+      "login": "AzulGarza",
+      "name": "azul",
       "avatar_url": "https://avatars.githubusercontent.com/u/10517170?v=4",
-      "profile": "https://github.com/FedericoGarza",
+      "profile": "https://github.com/AzulGarza",
       "contributions": [
         "code",
         "maintenance"

diff --git a/.github/workflows/build-docs.yaml b/.github/workflows/build-docs.yaml
@@ -1,7 +1,7 @@
 name: "build-docs"
 on:
-  push:
-    branches: ["main"]
+  release:
+    types: [released]
   pull_request:
     branches: ["main"]
   workflow_dispatch:
@@ -15,16 +15,16 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Clone repo
-        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
 
       - name: Clone docs repo
-        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           repository: Nixtla/docs
           ref: scripts
           path: docs-scripts
 
-      - uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # v5.1.1
+      - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
         with:
           python-version: '3.10'
           cache-dependency-path: settings.ini
@@ -46,7 +46,9 @@ jobs:
         run: cp nbs/mint.json _docs/mint.json && cp docs-scripts/imgs/* _docs/
 
       - name: Deploy to Mintlify Docs
-        if: github.event_name == 'push'
+        if: | 
+          github.event_name == 'release' || 
+          github.event_name == 'workflow_dispatch'
         uses: peaceiris/actions-gh-pages@4f9cc6602d3f66b9c108549d475ec49e8ef4d45e # v4.0.0
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
@@ -56,7 +58,9 @@ jobs:
           user_email: 41898282+github-actions[bot]@users.noreply.github.com
 
       - name: Trigger mintlify workflow
-        if: github.event_name == 'push'
+        if: | 
+          github.event_name == 'release' || 
+          github.event_name == 'workflow_dispatch'
         uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
         with:
           github-token: ${{ secrets.DOCS_WORKFLOW_TOKEN }}
@@ -72,7 +76,9 @@ jobs:
         run: python docs-scripts/configure-redirects.py neuralforecast
 
       - name: Deploy to Github Pages
-        if: github.event_name == 'push'
+        if: | 
+          github.event_name == 'release' || 
+          github.event_name == 'workflow_dispatch'
         uses: peaceiris/actions-gh-pages@4f9cc6602d3f66b9c108549d475ec49e8ef4d45e # v4.0.0
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -6,10 +6,6 @@ on:
   pull_request:
     branches: [main]
 
-defaults:
-  run:
-    shell: bash -l {0}
-
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: true
@@ -21,26 +17,27 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-13, windows-latest]
-        python-version: ['3.8', '3.9', '3.10', '3.11']
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         exclude:
           - os: windows-latest
-            python-version: '3.11'
+            python-version: "3.11"
     env:
       AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID_NIXTLA_TMP }}
       AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY_NIXTLA_TMP }}
     steps:
       - name: Clone repo
-        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
 
       - name: Set up environment
-        uses: mamba-org/setup-micromamba@f8b8a1e23a26f60a44c853292711bacfd3eac822 # v1.9.0
+        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # 5.3.0
         with:
-          environment-file: environment-cpu.yml
-          create-args: python=${{ matrix.python-version }}
-          cache-environment: true
+          python-version: ${{ matrix.python-version }}
 
       - name: Install pip requirements
-        run: pip install ./
+        run: >
+          pip install uv &&
+          uv pip install --system -i https://download.pytorch.org/whl/cpu torch &&
+          uv pip install --system "numpy<2" ".[dev]"
 
       - name: Tests
         run: nbdev_test --do_print --timing --n_workers 0 --flags polars
diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml
@@ -11,10 +11,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Clone repo
-        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
 
       - name: Set up python
-        uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # 5.1.1
+        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # 5.3.0
         with:
           python-version: 3.8
 

diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
@@ -21,9 +21,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
       - name: Set up Python
-        uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # 5.1.1
+        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # 5.3.0
         with:
           python-version: '3.x'
       - name: Install dependencies
@@ -33,7 +33,7 @@ jobs:
       - name: Build package
         run: python -m build
       - name: Publish package
-        uses: pypa/gh-action-pypi-publish@ec4db0b4ddc65acdf4bff5fa45ac92d78b56bdf0 # v1.9.0
+        uses: pypa/gh-action-pypi-publish@15c56dba361d8335944d31a2ecd17d700fc7bcbc # v1.12.2
         with:
           user: __token__
           password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/README.md b/README.md
@@ -115,7 +115,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/d
 <table>
   <tbody>
     <tr>
-      <td align="center" valign="top" width="14.28%"><a href="https://github.com/AzulGarza"><img src="https://avatars.githubusercontent.com/u/10517170?v=4?s=100" width="100px;" alt="Azul"/><br /><sub><b>fede</b></sub></a><br /><a href="https://github.com/Nixtla/neuralforecast/commits?author=AzulGarza" title="Code">💻</a> <a href="#maintenance-AzulGarza" title="Maintenance">🚧</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/AzulGarza"><img src="https://avatars.githubusercontent.com/u/10517170?v=4?s=100" width="100px;" alt="azul"/><br /><sub><b>azul</b></sub></a><br /><a href="https://github.com/Nixtla/neuralforecast/commits?author=AzulGarza" title="Code">💻</a> <a href="#maintenance-AzulGarza" title="Maintenance">🚧</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/cchallu"><img src="https://avatars.githubusercontent.com/u/31133398?v=4?s=100" width="100px;" alt="Cristian Challu"/><br /><sub><b>Cristian Challu</b></sub></a><br /><a href="https://github.com/Nixtla/neuralforecast/commits?author=cchallu" title="Code">💻</a> <a href="#maintenance-cchallu" title="Maintenance">🚧</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/jmoralez"><img src="https://avatars.githubusercontent.com/u/8473587?v=4?s=100" width="100px;" alt="José Morales"/><br /><sub><b>José Morales</b></sub></a><br /><a href="https://github.com/Nixtla/neuralforecast/commits?author=jmoralez" title="Code">💻</a> <a href="#maintenance-jmoralez" title="Maintenance">🚧</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/mergenthaler"><img src="https://avatars.githubusercontent.com/u/4086186?v=4?s=100" width="100px;" alt="mergenthaler"/><br /><sub><b>mergenthaler</b></sub></a><br /><a href="https://github.com/Nixtla/neuralforecast/commits?author=mergenthaler" title="Documentation">📖</a> <a href="https://github.com/Nixtla/neuralforecast/commits?author=mergenthaler" title="Code">💻</a></td>

diff --git a/action_files/test_models/src/multivariate_models.py b/action_files/test_models/src/multivariate_models.py
@@ -10,7 +10,7 @@
 from neuralforecast.models.tsmixer import TSMixer
 from neuralforecast.models.tsmixerx import TSMixerx
 from neuralforecast.models.itransformer import iTransformer
-from neuralforecast.models.stemgnn import StemGNN
+# from neuralforecast.models.stemgnn import StemGNN
 from neuralforecast.models.mlpmultivariate import MLPMultivariate
 from neuralforecast.models.timemixer import TimeMixer
 
@@ -30,7 +30,7 @@ def main(dataset: str = 'multivariate', group: str = 'ETTm2') -> None:
         TSMixer(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500),
         TSMixerx(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500),
         iTransformer(h=horizon, n_series=7, input_size=2 * horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500),
-        StemGNN(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout_rate=0.0, max_steps=1000, val_check_steps=500),
+        # StemGNN(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout_rate=0.0, max_steps=1000, val_check_steps=500),
         MLPMultivariate(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), max_steps=1000, val_check_steps=500),
         TimeMixer(h=horizon, n_series=7, input_size=2*horizon, loss=MAE(), dropout=0.0, max_steps=1000, val_check_steps=500)
     ]

diff --git a/nbs/common.base_auto.ipynb b/nbs/common.base_auto.ipynb
@@ -85,8 +85,8 @@
     "class MockTrial:\n",
     "    def suggest_int(*args, **kwargs):\n",
     "        return 'int'\n",
-    "    def suggest_categorical(*args, **kwargs):\n",
-    "        return 'categorical'\n",
+    "    def suggest_categorical(self, name, choices):\n",
+    "        return choices\n",
     "    def suggest_uniform(*args, **kwargs):\n",
     "        return 'uniform'\n",
     "    def suggest_loguniform(*args, **kwargs):\n",

diff --git a/nbs/common.base_model.ipynb b/nbs/common.base_model.ipynb
@@ -360,13 +360,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",
@@ -487,7 +499,7 @@
     "        if self.val_size == 0:\n",
     "            return\n",
     "        losses = torch.stack(self.validation_step_outputs)\n",
-    "        avg_loss = losses.mean().item()\n",
+    "        avg_loss = losses.mean().detach().item()\n",
     "        self.log(\n",
     "            \"ptl/val_loss\",\n",
     "            avg_loss,\n",

diff --git a/nbs/common.base_multivariate.ipynb b/nbs/common.base_multivariate.ipynb
@@ -109,6 +109,7 @@
     "                 optimizer_kwargs=None,\n",
     "                 lr_scheduler=None,\n",
     "                 lr_scheduler_kwargs=None,\n",
+    "                 dataloader_kwargs=None,\n",
     "                 **trainer_kwargs):\n",
     "        super().__init__(\n",
     "            random_seed=random_seed,\n",
@@ -131,7 +132,7 @@
     "        self.h = h\n",
     "        self.input_size = input_size\n",
     "        self.n_series = n_series\n",
-    "        self.padder = nn.ConstantPad1d(padding=(0, self.h), value=0)\n",
+    "        self.padder = nn.ConstantPad1d(padding=(0, self.h), value=0.0)\n",
     "\n",
     "        # Multivariate models do not support these loss functions yet.\n",
     "        unsupported_losses = (\n",
@@ -173,6 +174,7 @@
     "\n",
     "        # DataModule arguments\n",
     "        self.num_workers_loader = num_workers_loader\n",
+    "        self.dataloader_kwargs = dataloader_kwargs\n",
     "        self.drop_last_loader = drop_last_loader\n",
     "        # used by on_validation_epoch_end hook\n",
     "        self.validation_step_outputs = []\n",
@@ -390,12 +392,12 @@
     "\n",
     "        self.log(\n",
     "            'train_loss',\n",
-    "            loss.item(),\n",
+    "            loss.detach().item(),\n",
     "            batch_size=outsample_y.size(0),\n",
     "            prog_bar=True,\n",
     "            on_epoch=True,\n",
     "        )\n",
-    "        self.train_trajectories.append((self.global_step, loss.item()))\n",
+    "        self.train_trajectories.append((self.global_step, loss.detach().item()))\n",
     "        return loss\n",
     "\n",
     "    def validation_step(self, batch, batch_idx):\n",
@@ -440,7 +442,7 @@
     "\n",
     "        self.log(\n",
     "            'valid_loss',\n",
-    "            valid_loss.item(),\n",
+    "            valid_loss.detach().item(),\n",
     "            batch_size=outsample_y.size(0),\n",
     "            prog_bar=True,\n",
     "            on_epoch=True,\n",

diff --git a/nbs/common.base_recurrent.ipynb b/nbs/common.base_recurrent.ipynb
@@ -115,6 +115,7 @@
     "                 optimizer_kwargs=None,\n",
     "                 lr_scheduler=None,\n",
     "                 lr_scheduler_kwargs=None,\n",
+    "                 dataloader_kwargs=None,\n",
     "                 **trainer_kwargs):\n",
     "        super().__init__(\n",
     "            random_seed=random_seed,\n",
@@ -137,7 +138,7 @@
     "        self.h = h\n",
     "        self.input_size = input_size\n",
     "        self.inference_input_size = inference_input_size\n",
-    "        self.padder = nn.ConstantPad1d(padding=(0, self.h), value=0)\n",
+    "        self.padder = nn.ConstantPad1d(padding=(0, self.h), value=0.0)\n",
     "\n",
     "        unsupported_distributions = ['Bernoulli', 'ISQF']\n",
     "        if isinstance(self.loss, losses.DistributionLoss) and\\\n",
@@ -172,6 +173,7 @@
     "\n",
     "        # DataModule arguments\n",
     "        self.num_workers_loader = num_workers_loader\n",
+    "        self.dataloader_kwargs = dataloader_kwargs\n",
     "        self.drop_last_loader = drop_last_loader\n",
     "        # used by on_validation_epoch_end hook\n",
     "        self.validation_step_outputs = []\n",
@@ -254,7 +256,7 @@
     "\n",
     "            # Test size covers all data, pad left one timestep with zeros\n",
     "            if temporal.shape[-1] == self.test_size:\n",
-    "                padder_left = nn.ConstantPad1d(padding=(1, 0), value=0)\n",
+    "                padder_left = nn.ConstantPad1d(padding=(1, 0), value=0.0)\n",
     "                temporal = padder_left(temporal)\n",
     "\n",
     "        # Parse batch\n",
@@ -365,12 +367,12 @@
     "\n",
     "        self.log(\n",
     "            'train_loss',\n",
-    "            loss.item(),\n",
+    "            loss.detach().item(),\n",
     "            batch_size=outsample_y.size(0),\n",
     "            prog_bar=True,\n",
     "            on_epoch=True,\n",
     "        )\n",
-    "        self.train_trajectories.append((self.global_step, loss.item()))\n",
+    "        self.train_trajectories.append((self.global_step, loss.detach().item()))\n",
     "        return loss\n",
     "\n",
     "    def validation_step(self, batch, batch_idx):\n",
@@ -438,7 +440,7 @@
     "\n",
     "        self.log(\n",
     "            'valid_loss',\n",
-    "            valid_loss.item(),\n",
+    "            valid_loss.detach().item(),\n",
     "            batch_size=outsample_y.size(0),\n",
     "            prog_bar=True,\n",
     "            on_epoch=True,\n",
@@ -536,7 +538,7 @@
     "        self._check_exog(dataset)\n",
     "        self._restart_seed(random_seed)\n",
     "        data_module_kwargs = self._set_quantile_for_iqloss(**data_module_kwargs)\n",
-    "\n",
+    "        \n",
     "        if step_size > 1:\n",
     "            raise Exception('Recurrent models do not support step_size > 1')\n",
     "\n",