diff --git a/examples/ptf_V2_example.ipynb b/examples/ptf_V2_example.ipynb new file mode 100644 index 000000000..71e90e9de --- /dev/null +++ b/examples/ptf_V2_example.ipynb @@ -0,0 +1,1018 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "2630DaOEI4AJ", + "outputId": "b4cc7100-2a9c-41a3-e890-164b90c91c03" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Collecting pytorch-forecasting\n", + " Downloading pytorch_forecasting-1.3.0-py3-none-any.whl.metadata (13 kB)\n", + "Requirement already satisfied: numpy<=3.0.0 in /usr/local/lib/python3.11/dist-packages (from pytorch-forecasting) (2.0.2)\n", + "Requirement already satisfied: torch!=2.0.1,<3.0.0,>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from pytorch-forecasting) (2.6.0+cu124)\n", + "Collecting lightning<3.0.0,>=2.0.0 (from pytorch-forecasting)\n", + " Downloading lightning-2.5.1.post0-py3-none-any.whl.metadata (39 kB)\n", + "Requirement already satisfied: scipy<2.0,>=1.8 in /usr/local/lib/python3.11/dist-packages (from pytorch-forecasting) (1.15.2)\n", + "Requirement already satisfied: pandas<3.0.0,>=1.3.0 in /usr/local/lib/python3.11/dist-packages (from pytorch-forecasting) (2.2.2)\n", + "Requirement already satisfied: scikit-learn<2.0,>=1.2 in /usr/local/lib/python3.11/dist-packages (from pytorch-forecasting) (1.6.1)\n", + "Requirement already satisfied: PyYAML<8.0,>=5.4 in /usr/local/lib/python3.11/dist-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (6.0.2)\n", + "Requirement already satisfied: fsspec<2026.0,>=2022.5.0 in /usr/local/lib/python3.11/dist-packages (from fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (2025.3.2)\n", + "Collecting lightning-utilities<2.0,>=0.10.0 (from lightning<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading lightning_utilities-0.14.3-py3-none-any.whl.metadata (5.6 kB)\n", + "Requirement already satisfied: packaging<25.0,>=20.0 in /usr/local/lib/python3.11/dist-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (24.2)\n", + "Collecting torchmetrics<3.0,>=0.7.0 (from lightning<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading torchmetrics-1.7.1-py3-none-any.whl.metadata (21 kB)\n", + "Requirement already satisfied: tqdm<6.0,>=4.57.0 in /usr/local/lib/python3.11/dist-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (4.67.1)\n", + "Requirement already satisfied: typing-extensions<6.0,>=4.4.0 in /usr/local/lib/python3.11/dist-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (4.13.2)\n", + "Collecting pytorch-lightning (from lightning<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading pytorch_lightning-2.5.1.post0-py3-none-any.whl.metadata (20 kB)\n", + "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/dist-packages (from pandas<3.0.0,>=1.3.0->pytorch-forecasting) (2.9.0.post0)\n", + "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas<3.0.0,>=1.3.0->pytorch-forecasting) (2025.2)\n", + "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas<3.0.0,>=1.3.0->pytorch-forecasting) (2025.2)\n", + "Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.11/dist-packages (from scikit-learn<2.0,>=1.2->pytorch-forecasting) (1.4.2)\n", + "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.11/dist-packages (from scikit-learn<2.0,>=1.2->pytorch-forecasting) (3.6.0)\n", + "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.18.0)\n", + "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.4.2)\n", + "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.1.6)\n", + "Collecting nvidia-cuda-nvrtc-cu12==12.4.127 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", + "Collecting nvidia-cuda-runtime-cu12==12.4.127 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", + "Collecting nvidia-cuda-cupti-cu12==12.4.127 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", + "Collecting nvidia-cudnn-cu12==9.1.0.70 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", + "Collecting nvidia-cublas-cu12==12.4.5.8 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", + "Collecting nvidia-cufft-cu12==11.2.1.3 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", + "Collecting nvidia-curand-cu12==10.3.5.147 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", + "Collecting nvidia-cusolver-cu12==11.6.1.9 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", + "Collecting nvidia-cusparse-cu12==12.3.1.170 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n", + "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.2 in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (0.6.2)\n", + "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (2.21.5)\n", + "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.4.127)\n", + "Collecting nvidia-nvjitlink-cu12==12.4.127 (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting)\n", + " Downloading nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n", + "Requirement already satisfied: triton==3.2.0 in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.2.0)\n", + "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (1.13.1)\n", + "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (1.3.0)\n", + "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.11/dist-packages (from fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (3.11.15)\n", + "Requirement already satisfied: setuptools in /usr/local/lib/python3.11/dist-packages (from lightning-utilities<2.0,>=0.10.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (75.2.0)\n", + "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.11/dist-packages (from python-dateutil>=2.8.2->pandas<3.0.0,>=1.3.0->pytorch-forecasting) (1.17.0)\n", + "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.0.2)\n", + "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (2.6.1)\n", + "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (1.3.2)\n", + "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (25.3.0)\n", + "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (1.6.0)\n", + "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (6.4.3)\n", + "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (0.3.1)\n", + "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (1.20.0)\n", + "Requirement already satisfied: idna>=2.0 in /usr/local/lib/python3.11/dist-packages (from yarl<2.0,>=1.17.0->aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (3.10)\n", + "Downloading pytorch_forecasting-1.3.0-py3-none-any.whl (197 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m197.7/197.7 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading lightning-2.5.1.post0-py3-none-any.whl (819 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m819.0/819.0 kB\u001b[0m \u001b[31m16.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl (363.4 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m363.4/363.4 MB\u001b[0m \u001b[31m2.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (13.8 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.8/13.8 MB\u001b[0m \u001b[31m14.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (24.6 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m24.6/24.6 MB\u001b[0m \u001b[31m37.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (883 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m883.7/883.7 kB\u001b[0m \u001b[31m20.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl (664.8 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m664.8/664.8 MB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl (211.5 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m211.5/211.5 MB\u001b[0m \u001b[31m5.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl (56.3 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m12.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl (127.9 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m127.9/127.9 MB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl (207.5 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.5/207.5 MB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (21.1 MB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m64.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading lightning_utilities-0.14.3-py3-none-any.whl (28 kB)\n", + "Downloading torchmetrics-1.7.1-py3-none-any.whl (961 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m961.5/961.5 kB\u001b[0m \u001b[31m35.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hDownloading pytorch_lightning-2.5.1.post0-py3-none-any.whl (823 kB)\n", + "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m823.1/823.1 kB\u001b[0m \u001b[31m37.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", + "\u001b[?25hInstalling collected packages: nvidia-nvjitlink-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, lightning-utilities, nvidia-cusparse-cu12, nvidia-cudnn-cu12, nvidia-cusolver-cu12, torchmetrics, pytorch-lightning, lightning, pytorch-forecasting\n", + " Attempting uninstall: nvidia-nvjitlink-cu12\n", + " Found existing installation: nvidia-nvjitlink-cu12 12.5.82\n", + " Uninstalling nvidia-nvjitlink-cu12-12.5.82:\n", + " Successfully uninstalled nvidia-nvjitlink-cu12-12.5.82\n", + " Attempting uninstall: nvidia-curand-cu12\n", + " Found existing installation: nvidia-curand-cu12 10.3.6.82\n", + " Uninstalling nvidia-curand-cu12-10.3.6.82:\n", + " Successfully uninstalled nvidia-curand-cu12-10.3.6.82\n", + " Attempting uninstall: nvidia-cufft-cu12\n", + " Found existing installation: nvidia-cufft-cu12 11.2.3.61\n", + " Uninstalling nvidia-cufft-cu12-11.2.3.61:\n", + " Successfully uninstalled nvidia-cufft-cu12-11.2.3.61\n", + " Attempting uninstall: nvidia-cuda-runtime-cu12\n", + " Found existing installation: nvidia-cuda-runtime-cu12 12.5.82\n", + " Uninstalling nvidia-cuda-runtime-cu12-12.5.82:\n", + " Successfully uninstalled nvidia-cuda-runtime-cu12-12.5.82\n", + " Attempting uninstall: nvidia-cuda-nvrtc-cu12\n", + " Found existing installation: nvidia-cuda-nvrtc-cu12 12.5.82\n", + " Uninstalling nvidia-cuda-nvrtc-cu12-12.5.82:\n", + " Successfully uninstalled nvidia-cuda-nvrtc-cu12-12.5.82\n", + " Attempting uninstall: nvidia-cuda-cupti-cu12\n", + " Found existing installation: nvidia-cuda-cupti-cu12 12.5.82\n", + " Uninstalling nvidia-cuda-cupti-cu12-12.5.82:\n", + " Successfully uninstalled nvidia-cuda-cupti-cu12-12.5.82\n", + " Attempting uninstall: nvidia-cublas-cu12\n", + " Found existing installation: nvidia-cublas-cu12 12.5.3.2\n", + " Uninstalling nvidia-cublas-cu12-12.5.3.2:\n", + " Successfully uninstalled nvidia-cublas-cu12-12.5.3.2\n", + " Attempting uninstall: nvidia-cusparse-cu12\n", + " Found existing installation: nvidia-cusparse-cu12 12.5.1.3\n", + " Uninstalling nvidia-cusparse-cu12-12.5.1.3:\n", + " Successfully uninstalled nvidia-cusparse-cu12-12.5.1.3\n", + " Attempting uninstall: nvidia-cudnn-cu12\n", + " Found existing installation: nvidia-cudnn-cu12 9.3.0.75\n", + " Uninstalling nvidia-cudnn-cu12-9.3.0.75:\n", + " Successfully uninstalled nvidia-cudnn-cu12-9.3.0.75\n", + " Attempting uninstall: nvidia-cusolver-cu12\n", + " Found existing installation: nvidia-cusolver-cu12 11.6.3.83\n", + " Uninstalling nvidia-cusolver-cu12-11.6.3.83:\n", + " Successfully uninstalled nvidia-cusolver-cu12-11.6.3.83\n", + "Successfully installed lightning-2.5.1.post0 lightning-utilities-0.14.3 nvidia-cublas-cu12-12.4.5.8 nvidia-cuda-cupti-cu12-12.4.127 nvidia-cuda-nvrtc-cu12-12.4.127 nvidia-cuda-runtime-cu12-12.4.127 nvidia-cudnn-cu12-9.1.0.70 nvidia-cufft-cu12-11.2.1.3 nvidia-curand-cu12-10.3.5.147 nvidia-cusolver-cu12-11.6.1.9 nvidia-cusparse-cu12-12.3.1.170 nvidia-nvjitlink-cu12-12.4.127 pytorch-forecasting-1.3.0 pytorch-lightning-2.5.1.post0 torchmetrics-1.7.1\n" + ] + } + ], + "source": [ + "%pip install pytorch-forecasting" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "M7PQerTbI_tM" + }, + "outputs": [], + "source": [ + "from lightning.pytorch import Trainer\n", + "import numpy as np\n", + "import pandas as pd\n", + "from sklearn.preprocessing import RobustScaler, StandardScaler\n", + "import torch\n", + "import torch.nn as nn\n", + "from torch.optim import Optimizer\n", + "from torch.utils.data import Dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "DGTyf3vct-Jk" + }, + "outputs": [], + "source": [ + "from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule\n", + "from pytorch_forecasting.data.encoders import (\n", + " EncoderNormalizer,\n", + " NaNLabelEncoder,\n", + " TorchNormalizer,\n", + ")\n", + "from pytorch_forecasting.data.timeseries import TimeSeries\n", + "from pytorch_forecasting.metrics import MAE, SMAPE\n", + "from pytorch_forecasting.models.temporal_fusion_transformer.tft_version_two import TFT" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 206 + }, + "id": "WX-FRdusJSVN", + "outputId": "d162e241-3076-415c-db39-8c571bbaa282" + }, + "outputs": [ + { + "data": { + "application/vnd.google.colaboratory.intrinsic+json": { + "summary": "{\n \"name\": \"data_df\",\n \"rows\": 4900,\n \"fields\": [\n {\n \"column\": \"series_id\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 28,\n \"min\": 0,\n \"max\": 99,\n \"num_unique_values\": 100,\n \"samples\": [\n 83,\n 53,\n 70\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n },\n {\n \"column\": \"time_idx\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 14,\n \"min\": 0,\n \"max\": 48,\n \"num_unique_values\": 49,\n \"samples\": [\n 13,\n 45,\n 47\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n },\n {\n \"column\": \"x\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 0.6702424716860947,\n \"min\": -1.2572875930191487,\n \"max\": 1.347291996576924,\n \"num_unique_values\": 4900,\n \"samples\": [\n 0.40064266948811306,\n 0.688757012378203,\n -0.9278241195910876\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n },\n {\n \"column\": \"y\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 0.6742204890063661,\n \"min\": -1.2572875930191487,\n \"max\": 1.347291996576924,\n \"num_unique_values\": 4900,\n \"samples\": [\n 0.6039073395571968,\n 0.5832480743546181,\n -0.801772762118357\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n },\n {\n \"column\": \"category\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 1,\n \"min\": 0,\n \"max\": 4,\n \"num_unique_values\": 5,\n \"samples\": [\n 1,\n 4,\n 2\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n },\n {\n \"column\": \"future_known_feature\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 0.6741140972121411,\n \"min\": -0.9991351502732795,\n \"max\": 1.0,\n \"num_unique_values\": 49,\n \"samples\": [\n 0.26749882862458735,\n -0.2107957994307797,\n -0.01238866346289056\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n },\n {\n \"column\": \"static_feature\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 0.2912918986931303,\n \"min\": 0.007584244652032224,\n \"max\": 0.9959799570401108,\n \"num_unique_values\": 100,\n \"samples\": [\n 0.4307103381570838,\n 0.6664272198589233,\n 0.16731443141739688\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n },\n {\n \"column\": \"static_feature_cat\",\n \"properties\": {\n \"dtype\": \"number\",\n \"std\": 0,\n \"min\": 0,\n \"max\": 2,\n \"num_unique_values\": 3,\n \"samples\": [\n 0,\n 1,\n 2\n ],\n \"semantic_type\": \"\",\n \"description\": \"\"\n }\n }\n ]\n}", + "type": "dataframe", + "variable_name": "data_df" + }, + "text/html": [ + "\n", + "
\n", + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
series_idtime_idxxycategoryfuture_known_featurestatic_featurestatic_feature_cat
0000.2009680.23229801.0000000.687290
1010.2322980.33666900.9950040.687290
2020.3366690.63606300.9800670.687290
3030.6360630.92771000.9553360.687290
4040.9277101.00855400.9210610.687290
\n", + "
\n", + "
\n", + "\n", + "
\n", + " \n", + "\n", + " \n", + "\n", + " \n", + "
\n", + "\n", + "\n", + "
\n", + " \n", + "\n", + "\n", + "\n", + " \n", + "
\n", + "\n", + "
\n", + "
\n" + ], + "text/plain": [ + " series_id time_idx x y category future_known_feature \\\n", + "0 0 0 0.200968 0.232298 0 1.000000 \n", + "1 0 1 0.232298 0.336669 0 0.995004 \n", + "2 0 2 0.336669 0.636063 0 0.980067 \n", + "3 0 3 0.636063 0.927710 0 0.955336 \n", + "4 0 4 0.927710 1.008554 0 0.921061 \n", + "\n", + " static_feature static_feature_cat \n", + "0 0.68729 0 \n", + "1 0.68729 0 \n", + "2 0.68729 0 \n", + "3 0.68729 0 \n", + "4 0.68729 0 " + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "num_series = 100\n", + "seq_length = 50\n", + "data_list = []\n", + "for i in range(num_series):\n", + " x = np.arange(seq_length)\n", + " y = np.sin(x / 5.0) + np.random.normal(scale=0.1, size=seq_length)\n", + " category = i % 5\n", + " static_value = np.random.rand()\n", + " for t in range(seq_length - 1):\n", + " data_list.append(\n", + " {\n", + " \"series_id\": i,\n", + " \"time_idx\": t,\n", + " \"x\": y[t],\n", + " \"y\": y[t + 1],\n", + " \"category\": category,\n", + " \"future_known_feature\": np.cos(t / 10),\n", + " \"static_feature\": static_value,\n", + " \"static_feature_cat\": i % 3,\n", + " }\n", + " )\n", + "data_df = pd.DataFrame(data_list)\n", + "data_df.head()" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "AxxPHK6AKSD2", + "outputId": "dd95173d-73c2-451b-8b67-c9cc7298cf9d" + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + ":106: UserWarning: TimeSeries is part of an experimental rework of the pytorch-forecasting data layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. For beta testing, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n", + " warnings.warn(\n" + ] + } + ], + "source": [ + "dataset = TimeSeries(\n", + " data=data_df,\n", + " time=\"time_idx\",\n", + " target=\"y\",\n", + " group=[\"series_id\"],\n", + " num=[\"x\", \"future_known_feature\", \"static_feature\"],\n", + " cat=[\"category\", \"static_feature_cat\"],\n", + " known=[\"future_known_feature\"],\n", + " unknown=[\"x\", \"category\"],\n", + " static=[\"static_feature\", \"static_feature_cat\"],\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": { + "id": "5U5Lr_ZFKX0s" + }, + "outputs": [], + "source": [ + "data_module = EncoderDecoderTimeSeriesDataModule(\n", + " time_series_dataset=dataset,\n", + " max_encoder_length=30,\n", + " max_prediction_length=1,\n", + " batch_size=32,\n", + " categorical_encoders={\n", + " \"category\": NaNLabelEncoder(add_nan=True),\n", + " \"static_feature_cat\": NaNLabelEncoder(add_nan=True),\n", + " },\n", + " scalers={\n", + " \"x\": StandardScaler(),\n", + " \"future_known_feature\": StandardScaler(),\n", + " \"static_feature\": StandardScaler(),\n", + " },\n", + " target_normalizer=TorchNormalizer(),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/", + "height": 1000, + "referenced_widgets": [ + "7de178cd43ab4104ba2445a057a5f1a4", + "40598800b1234eaeb769f18e3c27865c", + "05a21a6e18ca46e280a8c66d6a73cf81", + "e4a86f54cf33447f8959864533cbae72", + "d5ba9114135f4ec1818790a62beb865a", + "39ab0f799ae64a8c9c0d6041b17c0ba5", + "6b98518240b744fe8a708b37a2dcaabe", + "274c42415f834e679f45d02d7d1c01d7", + "c88fdff0012c4e4e85563adb36287774", + "c2490a4409f34a608f6073ff6b9426eb", + "1a5e8e6d619740c9b0fd752a8f886b0c", + "ceff330da01b4ed39eafa8820cb6a5ca", + "d1d794dcb83746a48628280e7d552a70", + "31e24662df144d7196e01b873ffed137", + "07c7871be3274f45a93f01c6003b35fb", + "86ab236d48d04f1980836770cfe61b0d", + "73b46d96f2d54ea8b132b409b0739588", + "0b8814890f1d4143842acce4df31d93e", + "7960939d00d844f094e1bcdc1acda7f9", + "2dcacb2bf93c4d16bfb8657670499fc5", + "ed187855e406486ca0ea2259f3e2f43c", + "c73f370b888f4893be08d51b08b23a87", + "ea890982ba5a4c3d8e15e6bbd7285f6a", + "02f4487acf82401b972b4593da15de15", + "cca521da65e946abad580e9db9f2ac6b", + "b4c9075a5c1148ee8e422b2fcf86d90b", + "57dd80719a2e444cab8345bf5086a2f7", + "14e2d31e04e9447ba4f0c55000e0abe2", + "b35a713d4d7041109db487d10c55aaed", + "59ca144ddc884bb5a1c038e96cdf0dc0", + "95f9b03446af41fb83876f625aae5d75", + "bd2154cf3b04468b93a3bec23b5e34fa", + "f4713a710f274dd389691d7c12a7e740", + "07571714d67e4b8793ee76b0fe151e67", + "4e7273caa91147019971bf75ffe71e49", + "b751dbe4e93341eaa7d1a7683c277d83", + "11e1c349cb894caa9fd77333f7ababb9", + "8c9ac67ac6af488fb16e64c834634a30", + "704f28911d674088b8dfc240c6e28449", + "67da733b7a254020ad4d8d3877eb4494", + "72858abca77d430a8d009ad72127b331", + "86e82ab383a54ef7aa7c0abec6faf1be", + "60f097a304044e3db6dda60ff381776a", + "334e7a44226f4414a825e2d81e9571c1", + "13246130b24145d680092a5a3929546e", + "e64a93136fbf468b8e503cd202dbc986", + "559df10442ac42fc8033bdf014864334", + "67e6440ba0424db18646d47beff2e37b", + "36707c56fcfb4fadbfff37d44ee52d5b", + "484702bf9d854cb6a964de47d6975aa6", + "ddb67e4d1749424299a2c40b36810809", + "6f1e8c2aa06a48548a7750869d3f056e", + "d7df1875fbbb4c78909d76c5b81ffd95", + "cb60f83efc234cd6a90212125fe841a1", + "dfca7d9f0f7844ea953ce4b659493695", + "14ed4e565d5a4b319b0c38557a935b92", + "48e7f65e9dcc48289f44724da248875d", + "8220ff572b31472e91dc5d7553ca43b9", + "fcd1b8b77a7f4c108ba02de9779f5bb2", + "a53e4183893242ea884db9f25e439d94", + "e54fdc83afdf492a9540bc892bdc262a", + "04ce55b4c81a46b09d55b8a0dc3fae00", + "86eddf636c2b45fc93a69f3c8a260b1c", + "287ce5b9bafb402e9698f20779f52386", + "ef80c0b7306046738cfd4e16f5af1afc", + "c0bb5f38ec9346119441c6cb8fd71c5a", + "a54b4228eaf34dc1b40ee8d40500e069", + "28cb77e3b48a4c669dcaa79609d332c2", + "8acc35455cb84ece9e7b3d84d8870a7a", + "afe9a01608f449479cc491f75e095d55", + "9b5971f9e8d44e8a872b782ef49d306a", + "1538291371bc45b2bcb7a687c6b8f79c", + "d595a038c20e4d349e392ff1795dd418", + "a67fa2461a18428892e57790074fd5b6", + "6733a54e084a4d348a85e574985720b7", + "204594abf12e43d4abf0fa35544fb64a", + "eee2d2ef09294652bd94ff768ffb99f0", + "c8b8814037464155935670b776378ab9", + "0b6b787693af4b548ad59ad7ba6c921f", + "6b7172a2b1fd4ddea3dba3d435bf36fe", + "ed0ffcd137e94edebad5e41956ad7466", + "0344079173a84c448cd6edbac612d970", + "890f17882dc44e76be06dd88e6d07cff", + "1c552ec98dcd4325800ee8c9dddc398a", + "89d54947a3d146e4b089a22a2711a234", + "91e7ef398d1646bcbfb598480d949c74", + "13173089480c49b18c4ba016da69a855", + "6b717d2759ca445da7642a2cf20f22b5" + ] + }, + "id": "Si7bbZIULBZz", + "outputId": "0b2f26b3-e37c-4ab6-c234-90693745a8cd" + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "INFO: Using default `ModelCheckpoint`. Consider installing `litmodels` package to enable `LitModelCheckpoint` for automatic upload to the Lightning model registry.\n", + "INFO:lightning.pytorch.utilities.rank_zero:Using default `ModelCheckpoint`. Consider installing `litmodels` package to enable `LitModelCheckpoint` for automatic upload to the Lightning model registry.\n", + "INFO: GPU available: False, used: False\n", + "INFO:lightning.pytorch.utilities.rank_zero:GPU available: False, used: False\n", + "INFO: TPU available: False, using: 0 TPU cores\n", + "INFO:lightning.pytorch.utilities.rank_zero:TPU available: False, using: 0 TPU cores\n", + "INFO: HPU available: False, using: 0 HPUs\n", + "INFO:lightning.pytorch.utilities.rank_zero:HPU available: False, using: 0 HPUs\n", + "INFO: \n", + " | Name | Type | Params | Mode \n", + "---------------------------------------------------------------------\n", + "0 | loss | MSELoss | 0 | train\n", + "1 | encoder_var_selection | Sequential | 709 | train\n", + "2 | decoder_var_selection | Sequential | 193 | train\n", + "3 | static_context_linear | Linear | 192 | train\n", + "4 | lstm_encoder | LSTM | 51.5 K | train\n", + "5 | lstm_decoder | LSTM | 50.4 K | train\n", + "6 | self_attention | MultiheadAttention | 16.6 K | train\n", + "7 | pre_output | Linear | 4.2 K | train\n", + "8 | output_layer | Linear | 65 | train\n", + "---------------------------------------------------------------------\n", + "123 K Trainable params\n", + "0 Non-trainable params\n", + "123 K Total params\n", + "0.495 Total estimated model params size (MB)\n", + "18 Modules in train mode\n", + "0 Modules in eval mode\n", + "INFO:lightning.pytorch.callbacks.model_summary:\n", + " | Name | Type | Params | Mode \n", + "---------------------------------------------------------------------\n", + "0 | loss | MSELoss | 0 | train\n", + "1 | encoder_var_selection | Sequential | 709 | train\n", + "2 | decoder_var_selection | Sequential | 193 | train\n", + "3 | static_context_linear | Linear | 192 | train\n", + "4 | lstm_encoder | LSTM | 51.5 K | train\n", + "5 | lstm_decoder | LSTM | 50.4 K | train\n", + "6 | self_attention | MultiheadAttention | 16.6 K | train\n", + "7 | pre_output | Linear | 4.2 K | train\n", + "8 | output_layer | Linear | 65 | train\n", + "---------------------------------------------------------------------\n", + "123 K Trainable params\n", + "0 Non-trainable params\n", + "123 K Total params\n", + "0.495 Total estimated model params size (MB)\n", + "18 Modules in train mode\n", + "0 Modules in eval mode\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Training model...\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "7de178cd43ab4104ba2445a057a5f1a4", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Sanity Checking: | | 0/? [00:00┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", + "┃ Test metric ┃ DataLoader 0 ┃\n", + "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", + "│ test_MAE │ 0.45287469029426575 │\n", + "│ test_SMAPE │ 0.942494809627533 │\n", + "│ test_loss │ 0.01396977063268423 │\n", + "└───────────────────────────┴───────────────────────────┘\n", + "\n" + ], + "text/plain": [ + "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", + "┃\u001b[1m \u001b[0m\u001b[1m Test metric \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m DataLoader 0 \u001b[0m\u001b[1m \u001b[0m┃\n", + "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", + "│\u001b[36m \u001b[0m\u001b[36m test_MAE \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 0.45287469029426575 \u001b[0m\u001b[35m \u001b[0m│\n", + "│\u001b[36m \u001b[0m\u001b[36m test_SMAPE \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 0.942494809627533 \u001b[0m\u001b[35m \u001b[0m│\n", + "│\u001b[36m \u001b[0m\u001b[36m test_loss \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 0.01396977063268423 \u001b[0m\u001b[35m \u001b[0m│\n", + "└───────────────────────────┴───────────────────────────┘\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Prediction shape: torch.Size([32, 1, 1])\n", + "First prediction values: [[0.11045122]]\n", + "First true values: [[-0.0491814]]\n", + "\n", + "TFT model test complete!\n" + ] + } + ], + "source": [ + "model = TFT(\n", + " loss=nn.MSELoss(),\n", + " logging_metrics=[MAE(), SMAPE()],\n", + " optimizer=\"adam\",\n", + " optimizer_params={\"lr\": 1e-3},\n", + " lr_scheduler=\"reduce_lr_on_plateau\",\n", + " lr_scheduler_params={\"mode\": \"min\", \"factor\": 0.1, \"patience\": 10},\n", + " hidden_size=64,\n", + " num_layers=2,\n", + " attention_head_size=4,\n", + " dropout=0.1,\n", + " metadata=data_module.metadata,\n", + ")\n", + "\n", + "print(\"\\nTraining model...\")\n", + "trainer = Trainer(\n", + " max_epochs=5,\n", + " accelerator=\"auto\",\n", + " devices=1,\n", + " enable_progress_bar=True,\n", + " log_every_n_steps=10,\n", + ")\n", + "\n", + "trainer.fit(model, data_module)\n", + "\n", + "print(\"\\nEvaluating model...\")\n", + "test_metrics = trainer.test(model, data_module)\n", + "\n", + "model.eval()\n", + "with torch.no_grad():\n", + " test_batch = next(iter(data_module.test_dataloader()))\n", + " x_test, y_test = test_batch\n", + " y_pred = model(x_test)\n", + "\n", + " print(\"\\nPrediction shape:\", y_pred[\"prediction\"].shape)\n", + " print(\"First prediction values:\", y_pred[\"prediction\"][0].cpu().numpy())\n", + " print(\"First true values:\", y_test[0].cpu().numpy())\n", + "print(\"\\nTFT model test complete!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "zVRwi2MvLGgc" + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/examples/tslib_v2_example.ipynb b/examples/tslib_v2_example.ipynb new file mode 100644 index 000000000..d68491d94 --- /dev/null +++ b/examples/tslib_v2_example.ipynb @@ -0,0 +1,1677 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c20f63be", + "metadata": {}, + "source": [ + "# NOTE\n", + "\n", + "This notebook is just an example to demonstrate D1 compatibility with the `TimeXer` model. Considering that there is no concrete design for a TSLib specific D2 layer, for the time being we are using the `EncoderDecoderDataModule` and `BaseModel` for implementing `TimeXer`. The implementation is rather confusing with many overlapping bits because the D2 isn't solely built for TSLib, but is a demonstration of how `TimeXer` works with the latest version of v2 rework and shows promise for more models from TSlib." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "7563b0a7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Requirement already satisfied: pytorch-forecasting in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (1.2.0)\n", + "Requirement already satisfied: numpy<=3.0.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pytorch-forecasting) (2.2.2)\n", + "Requirement already satisfied: torch!=2.0.1,<3.0.0,>=2.0.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pytorch-forecasting) (2.5.1)\n", + "Requirement already satisfied: lightning<3.0.0,>=2.0.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pytorch-forecasting) (2.5.0.post0)\n", + "Requirement already satisfied: scipy<2.0,>=1.8 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pytorch-forecasting) (1.15.1)\n", + "Requirement already satisfied: pandas<3.0.0,>=1.3.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pytorch-forecasting) (2.2.3)\n", + "Requirement already satisfied: scikit-learn<2.0,>=1.2 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pytorch-forecasting) (1.6.1)\n", + "Requirement already satisfied: PyYAML<8.0,>=5.4 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (6.0.2)\n", + "Requirement already satisfied: fsspec[http]<2026.0,>=2022.5.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (2024.12.0)\n", + "Requirement already satisfied: lightning-utilities<2.0,>=0.10.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (0.11.9)\n", + "Requirement already satisfied: packaging<25.0,>=20.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (24.2)\n", + "Requirement already satisfied: torchmetrics<3.0,>=0.7.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (1.6.1)\n", + "Requirement already satisfied: tqdm<6.0,>=4.57.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (4.67.1)\n", + "Requirement already satisfied: typing-extensions<6.0,>=4.4.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (4.12.2)\n", + "Requirement already satisfied: pytorch-lightning in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from lightning<3.0.0,>=2.0.0->pytorch-forecasting) (2.5.0.post0)\n", + "Requirement already satisfied: python-dateutil>=2.8.2 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pandas<3.0.0,>=1.3.0->pytorch-forecasting) (2.9.0.post0)\n", + "Requirement already satisfied: pytz>=2020.1 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pandas<3.0.0,>=1.3.0->pytorch-forecasting) (2024.2)\n", + "Requirement already satisfied: tzdata>=2022.7 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from pandas<3.0.0,>=1.3.0->pytorch-forecasting) (2025.1)\n", + "Requirement already satisfied: joblib>=1.2.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from scikit-learn<2.0,>=1.2->pytorch-forecasting) (1.4.2)\n", + "Requirement already satisfied: threadpoolctl>=3.1.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from scikit-learn<2.0,>=1.2->pytorch-forecasting) (3.5.0)\n", + "Requirement already satisfied: filelock in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.17.0)\n", + "Requirement already satisfied: networkx in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.4.2)\n", + "Requirement already satisfied: jinja2 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.1.5)\n", + "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.4.127)\n", + "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.4.127)\n", + "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.4.127)\n", + "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (9.1.0.70)\n", + "Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.4.5.8)\n", + "Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (11.2.1.3)\n", + "Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (10.3.5.147)\n", + "Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (11.6.1.9)\n", + "Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.3.1.170)\n", + "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (2.21.5)\n", + "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.4.127)\n", + "Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (12.4.127)\n", + "Requirement already satisfied: triton==3.1.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.1.0)\n", + "Requirement already satisfied: setuptools in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (75.8.0)\n", + "Requirement already satisfied: sympy==1.13.1 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (1.13.1)\n", + "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from sympy==1.13.1->torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (1.3.0)\n", + "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (3.11.11)\n", + "Requirement already satisfied: six>=1.5 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from python-dateutil>=2.8.2->pandas<3.0.0,>=1.3.0->pytorch-forecasting) (1.17.0)\n", + "Requirement already satisfied: MarkupSafe>=2.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from jinja2->torch!=2.0.1,<3.0.0,>=2.0.0->pytorch-forecasting) (3.0.2)\n", + "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (2.4.4)\n", + "Requirement already satisfied: aiosignal>=1.1.2 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (1.3.2)\n", + "Requirement already satisfied: attrs>=17.3.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (25.1.0)\n", + "Requirement already satisfied: frozenlist>=1.1.1 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (1.5.0)\n", + "Requirement already satisfied: multidict<7.0,>=4.5 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (6.1.0)\n", + "Requirement already satisfied: propcache>=0.2.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (0.2.1)\n", + "Requirement already satisfied: yarl<2.0,>=1.17.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (1.18.3)\n", + "Requirement already satisfied: idna>=2.0 in /home/pranav/Desktop/code/pytorch-forecasting/.venv/lib/python3.12/site-packages (from yarl<2.0,>=1.17.0->aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2026.0,>=2022.5.0->lightning<3.0.0,>=2.0.0->pytorch-forecasting) (3.10)\n", + "\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.2.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.1.1\u001b[0m\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n", + "Note: you may need to restart the kernel to use updated packages.\n" + ] + } + ], + "source": [ + "%pip install pytorch-forecasting\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "524fb344", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Any, Optional, Union\n", + "\n", + "from lightning.pytorch import Trainer\n", + "import numpy as np\n", + "import pandas as pd\n", + "from sklearn.preprocessing import RobustScaler, StandardScaler\n", + "import torch\n", + "import torch.nn as nn\n", + "from torch.optim import Optimizer\n", + "from torch.utils.data import Dataset" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "3f1b0019", + "metadata": {}, + "outputs": [], + "source": [ + "from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule\n", + "from pytorch_forecasting.data.encoders import (\n", + " EncoderNormalizer,\n", + " NaNLabelEncoder,\n", + " TorchNormalizer,\n", + ")\n", + "from pytorch_forecasting.data.timeseries import TimeSeries\n", + "from pytorch_forecasting.metrics import MAE, SMAPE\n", + "from pytorch_forecasting.models.timexer._timexer import TimeXer" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "453abaa2", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.microsoft.datawrangler.viewer.v0+json": { + "columns": [ + { + "name": "index", + "rawType": "int64", + "type": "integer" + }, + { + "name": "series_id", + "rawType": "int64", + "type": "integer" + }, + { + "name": "time_idx", + "rawType": "int64", + "type": "integer" + }, + { + "name": "x", + "rawType": "float64", + "type": "float" + }, + { + "name": "y", + "rawType": "float64", + "type": "float" + }, + { + "name": "static_feature", + "rawType": "float64", + "type": "float" + } + ], + "conversionMethod": "pd.DataFrame", + "ref": "1eaf488b-f9bd-44a2-be5d-6da4307bec55", + "rows": [ + [ + "0", + "0", + "0", + "-0.0149095238786125", + "0.060817338247073804", + "0.3881772925541149" + ], + [ + "1", + "0", + "1", + "0.060817338247073804", + "0.36381631980778206", + "0.3881772925541149" + ], + [ + "2", + "0", + "2", + "0.36381631980778206", + "0.5852826362835459", + "0.3881772925541149" + ], + [ + "3", + "0", + "3", + "0.5852826362835459", + "0.6755320182905834", + "0.3881772925541149" + ], + [ + "4", + "0", + "4", + "0.6755320182905834", + "0.7312102935054476", + "0.3881772925541149" + ] + ], + "shape": { + "columns": 5, + "rows": 5 + } + }, + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
series_idtime_idxxystatic_feature
000-0.0149100.0608170.388177
1010.0608170.3638160.388177
2020.3638160.5852830.388177
3030.5852830.6755320.388177
4040.6755320.7312100.388177
\n", + "
" + ], + "text/plain": [ + " series_id time_idx x y static_feature\n", + "0 0 0 -0.014910 0.060817 0.388177\n", + "1 0 1 0.060817 0.363816 0.388177\n", + "2 0 2 0.363816 0.585283 0.388177\n", + "3 0 3 0.585283 0.675532 0.388177\n", + "4 0 4 0.675532 0.731210 0.388177" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "num_series = 100\n", + "seq_length = 50\n", + "data_list = []\n", + "for i in range(num_series):\n", + " x = np.arange(seq_length)\n", + " y = np.sin(x / 5.0) + np.random.normal(scale=0.1, size=seq_length)\n", + " category = i % 5\n", + " static_value = np.random.rand()\n", + " for t in range(seq_length - 1):\n", + " data_list.append(\n", + " {\n", + " \"series_id\": i,\n", + " \"time_idx\": t,\n", + " \"x\": y[t],\n", + " \"y\": y[t + 1],\n", + " \"static_feature\": static_value,\n", + " }\n", + " )\n", + "data_df = pd.DataFrame(data_list)\n", + "data_df.head()" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ee0f975b", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/pranav/Desktop/code/pytorch-forecasting/pytorch_forecasting/data/timeseries/_timeseries_v2.py:106: UserWarning: TimeSeries is part of an experimental rework of the pytorch-forecasting data layer, scheduled for release with v2.0.0. The API is not stable and may change without prior warning. For beta testing, but not for stable production use. Feedback and suggestions are very welcome in pytorch-forecasting issue 1736, https://github.com/sktime/pytorch-forecasting/issues/1736\n", + " warnings.warn(\n" + ] + } + ], + "source": [ + "dataset = TimeSeries(\n", + " data=data_df,\n", + " time=\"time_idx\",\n", + " target=\"y\",\n", + " group=[\"series_id\"],\n", + " num=[\"x\", \"static_feature\"],\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "4f13b58f", + "metadata": {}, + "outputs": [], + "source": [ + "data_module = EncoderDecoderTimeSeriesDataModule(\n", + " time_series_dataset=dataset,\n", + " max_encoder_length=30,\n", + " max_prediction_length=1,\n", + " batch_size=16,\n", + " scalers={\"x\": StandardScaler(), \"static_feature\": StandardScaler()},\n", + " target_normalizer=TorchNormalizer(),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "33eb6e78", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'encoder_cat': 0,\n", + " 'encoder_cont': 2,\n", + " 'decoder_cat': 0,\n", + " 'decoder_cont': 0,\n", + " 'target': 1,\n", + " 'static_categorical_features': 0,\n", + " 'static_continuous_features': 0,\n", + " 'max_encoder_length': 30,\n", + " 'max_prediction_length': 1,\n", + " 'min_encoder_length': 30,\n", + " 'min_prediction_length': 1}" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "data_module.setup(stage=\"fit\")\n", + "data_module.metadata" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "7c943cd2", + "metadata": {}, + "outputs": [], + "source": [ + "model = TimeXer(\n", + " loss=nn.L1Loss(),\n", + " logging_metrics=[MAE(), SMAPE()],\n", + " context_length=30,\n", + " prediction_length=1,\n", + " task_name=\"long_term_forecast\",\n", + " features=\"MS\",\n", + " d_model=32,\n", + " n_heads=2,\n", + " e_layers=1,\n", + " d_ff=64,\n", + " dropout=0.1,\n", + " patch_length=1,\n", + " use_norm=False,\n", + " metadata=data_module.metadata,\n", + " optimizer=\"adam\",\n", + " optimizer_params={\"lr\": 1e-3},\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "a472a9b5", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "GPU available: True (cuda), used: True\n", + "TPU available: False, using: 0 TPU cores\n", + "HPU available: False, using: 0 HPUs\n" + ] + } + ], + "source": [ + "trainer = Trainer(\n", + " max_epochs=5,\n", + " accelerator=\"auto\",\n", + " devices=1,\n", + " enable_progress_bar=True,\n", + " log_every_n_steps=10,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "9ee9aa67", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "You are using a CUDA device ('NVIDIA GeForce RTX 3050 6GB Laptop GPU') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n", + "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n", + "\n", + " | Name | Type | Params | Mode \n", + "----------------------------------------------------------------\n", + "0 | loss | L1Loss | 0 | train\n", + "1 | en_embedding | EnEmbedding | 64 | train\n", + "2 | ex_embedding | DataEmbedding_inverted | 992 | train\n", + "3 | encoder | Encoder | 12.9 K | train\n", + "4 | head | FlattenHead | 993 | train\n", + "----------------------------------------------------------------\n", + "14.9 K Trainable params\n", + "0 Non-trainable params\n", + "14.9 K Total params\n", + "0.060 Total estimated model params size (MB)\n", + "36 Modules in train mode\n", + "0 Modules in eval mode\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "3913efebc2214dc489de8e1ff608c2f7", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Sanity Checking: | | 0/? [00:00 int: """Return number of time series in the dataset.""" @@ -216,69 +195,54 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]: weights : torch.Tensor of shape (n_timepoints,), optional Only included if weights are not `None`. """ - time = self.time - feature_cols = self.feature_cols - _target = self._target - _known = self._known - _static = self._static - _group = self._group - _groups = self._groups - _group_ids = self._group_ids - weight = self.weight - data_future = self.data_future - - group_id = _group_ids[index] - - if _group: - mask = _groups[group_id] + group_id = self._group_ids[index] + + if self.group: + mask = self._groups[group_id] data = self.data.loc[mask] else: data = self.data - cutoff_time = data[time].max() - - data_vals = data[time].values - data_tgt_vals = data[_target].values - data_feat_vals = data[feature_cols].values + cutoff_time = data[self.time].max() result = { - "t": data_vals, - "y": torch.tensor(data_tgt_vals), - "x": torch.tensor(data_feat_vals), + "t": data[self.time].values, + "y": torch.tensor(data[self.target].values), + "x": torch.tensor(data[self.feature_cols].values), "group": torch.tensor([hash(str(group_id))]), - "st": torch.tensor(data[_static].iloc[0].values if _static else []), + "st": torch.tensor(data[self.static].iloc[0].values if self.static else []), "cutoff_time": cutoff_time, } - if data_future is not None: - if _group: - future_mask = self.data_future.groupby(_group).groups[group_id] + if self.data_future is not None: + if self.group: + future_mask = self.data_future.groupby(self.group).groups[group_id] future_data = self.data_future.loc[future_mask] else: future_data = self.data_future - data_fut_vals = future_data[time].values - - combined_times = np.concatenate([data_vals, data_fut_vals]) + combined_times = np.concatenate( + [data[self.time].values, future_data[self.time].values] + ) combined_times = np.unique(combined_times) combined_times.sort() num_timepoints = len(combined_times) - x_merged = np.full((num_timepoints, len(feature_cols)), np.nan) - y_merged = np.full((num_timepoints, len(_target)), np.nan) + x_merged = np.full((num_timepoints, len(self.feature_cols)), np.nan) + y_merged = np.full((num_timepoints, len(self.target)), np.nan) current_time_indices = {t: i for i, t in enumerate(combined_times)} - for i, t in enumerate(data_vals): + for i, t in enumerate(data[self.time].values): idx = current_time_indices[t] - x_merged[idx] = data_feat_vals[i] - y_merged[idx] = data_tgt_vals[i] + x_merged[idx] = data[self.feature_cols].values[i] + y_merged[idx] = data[self.target].values[i] - for i, t in enumerate(data_fut_vals): + for i, t in enumerate(future_data[self.time].values): if t in current_time_indices: idx = current_time_indices[t] - for j, col in enumerate(_known): - if col in feature_cols: - feature_idx = feature_cols.index(col) + for j, col in enumerate(self.known): + if col in self.feature_cols: + feature_idx = self.feature_cols.index(col) x_merged[idx, feature_idx] = future_data[col].values[i] result.update( @@ -289,17 +253,17 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]: } ) - if weight: + if self.weight: if self.data_future is not None and self.weight in self.data_future.columns: weights_merged = np.full(num_timepoints, np.nan) - for i, t in enumerate(data_vals): + for i, t in enumerate(data[self.time].values): idx = current_time_indices[t] - weights_merged[idx] = data[weight].values[i] + weights_merged[idx] = data[self.weight].values[i] - for i, t in enumerate(data_fut_vals): + for i, t in enumerate(future_data[self.time].values): if t in current_time_indices and self.weight in future_data.columns: idx = current_time_indices[t] - weights_merged[idx] = future_data[weight].values[i] + weights_merged[idx] = future_data[self.weight].values[i] result["weights"] = torch.tensor(weights_merged, dtype=torch.float32) else: diff --git a/pytorch_forecasting/models/__init__.py b/pytorch_forecasting/models/__init__.py index 29aeb24f5..0a9d600f8 100644 --- a/pytorch_forecasting/models/__init__.py +++ b/pytorch_forecasting/models/__init__.py @@ -19,6 +19,7 @@ TemporalFusionTransformer, ) from pytorch_forecasting.models.tide import TiDEModel +from pytorch_forecasting.models.timexer import TimeXer __all__ = [ "NBeats", @@ -37,4 +38,5 @@ "MultiEmbedding", "DecoderMLP", "TiDEModel", + "TimeXer", ] diff --git a/pytorch_forecasting/models/base/base_model_refactor.py b/pytorch_forecasting/models/base/base_model_refactor.py new file mode 100644 index 000000000..f03d70020 --- /dev/null +++ b/pytorch_forecasting/models/base/base_model_refactor.py @@ -0,0 +1,283 @@ +######################################################################################## +# Disclaimer: This baseclass is still work in progress and experimental, please +# use with care. This class is a basic skeleton of how the base classes may look like +# in the version-2. +######################################################################################## + + +from typing import Optional, Union + +from lightning.pytorch import LightningModule +from lightning.pytorch.utilities.types import STEP_OUTPUT +import torch +import torch.nn as nn +from torch.optim import Optimizer + + +class BaseModel(LightningModule): + def __init__( + self, + loss: nn.Module, + logging_metrics: Optional[list[nn.Module]] = None, + optimizer: Optional[Union[Optimizer, str]] = "adam", + optimizer_params: Optional[dict] = None, + lr_scheduler: Optional[str] = None, + lr_scheduler_params: Optional[dict] = None, + ): + """ + Base model for time series forecasting. + + Parameters + ---------- + loss : nn.Module + Loss function to use for training. + logging_metrics : Optional[list[nn.Module]], optional + list of metrics to log during training, validation, and testing. + optimizer : Optional[Union[Optimizer, str]], optional + Optimizer to use for training. + Can be a string ("adam", "sgd") or an instance of `torch.optim.Optimizer`. + optimizer_params : Optional[dict], optional + Parameters for the optimizer. + lr_scheduler : Optional[str], optional + Learning rate scheduler to use. + Supported values: "reduce_lr_on_plateau", "step_lr". + lr_scheduler_params : Optional[dict], optional + Parameters for the learning rate scheduler. + """ + super().__init__() + self.loss = loss + self.logging_metrics = logging_metrics if logging_metrics is not None else [] + self.optimizer = optimizer + self.optimizer_params = optimizer_params if optimizer_params is not None else {} + self.lr_scheduler = lr_scheduler + self.lr_scheduler_params = ( + lr_scheduler_params if lr_scheduler_params is not None else {} + ) + + def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Forward pass of the model. + + Parameters + ---------- + x : dict[str, torch.Tensor] + Dictionary containing input tensors + + Returns + ------- + dict[str, torch.Tensor] + Dictionary containing output tensors + """ + raise NotImplementedError("Forward method must be implemented by subclass.") + + def training_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> STEP_OUTPUT: + """ + Training step for the model. + + Parameters + ---------- + batch : tuple[dict[str, torch.Tensor]] + Batch of data containing input and target tensors. + batch_idx : int + Index of the batch. + + Returns + ------- + STEP_OUTPUT + Dictionary containing the loss and other metrics. + """ + x, y = batch + y_hat_dict = self(x) + y_hat = y_hat_dict["prediction"] + loss = self.loss(y_hat, y) + self.log( + "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True + ) + self.log_metrics(y_hat, y, prefix="train") + return {"loss": loss} + + def validation_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> STEP_OUTPUT: + """ + Validation step for the model. + + Parameters + ---------- + batch : tuple[dict[str, torch.Tensor]] + Batch of data containing input and target tensors. + batch_idx : int + Index of the batch. + + Returns + ------- + STEP_OUTPUT + Dictionary containing the loss and other metrics. + """ + x, y = batch + y_hat_dict = self(x) + y_hat = y_hat_dict["prediction"] + loss = self.loss(y_hat, y) + self.log( + "val_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True + ) + self.log_metrics(y_hat, y, prefix="val") + return {"val_loss": loss} + + def test_step( + self, batch: tuple[dict[str, torch.Tensor]], batch_idx: int + ) -> STEP_OUTPUT: + """ + Test step for the model. + + Parameters + ---------- + batch : tuple[dict[str, torch.Tensor]] + Batch of data containing input and target tensors. + batch_idx : int + Index of the batch. + + Returns + ------- + STEP_OUTPUT + Dictionary containing the loss and other metrics. + """ + x, y = batch + y_hat_dict = self(x) + y_hat = y_hat_dict["prediction"] + loss = self.loss(y_hat, y) + self.log( + "test_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True + ) + self.log_metrics(y_hat, y, prefix="test") + return {"test_loss": loss} + + def predict_step( + self, + batch: tuple[dict[str, torch.Tensor]], + batch_idx: int, + dataloader_idx: int = 0, + ) -> torch.Tensor: + """ + Prediction step for the model. + + Parameters + ---------- + batch : tuple[dict[str, torch.Tensor]] + Batch of data containing input tensors. + batch_idx : int + Index of the batch. + dataloader_idx : int + Index of the dataloader. + + Returns + ------- + torch.Tensor + Predicted output tensor. + """ + x, _ = batch + y_hat = self(x) + return y_hat + + def configure_optimizers(self) -> dict: + """ + Configure the optimizer and learning rate scheduler. + + Returns + ------- + dict + Dictionary containing the optimizer and scheduler configuration. + """ + optimizer = self._get_optimizer() + if self.lr_scheduler is not None: + scheduler = self._get_scheduler(optimizer) + if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + return { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": scheduler, + "monitor": "val_loss", + }, + } + else: + return {"optimizer": optimizer, "lr_scheduler": scheduler} + return {"optimizer": optimizer} + + def _get_optimizer(self) -> Optimizer: + """ + Get the optimizer based on the specified optimizer name and parameters. + + Returns + ------- + Optimizer + The optimizer instance. + """ + if isinstance(self.optimizer, str): + if self.optimizer.lower() == "adam": + return torch.optim.Adam(self.parameters(), **self.optimizer_params) + elif self.optimizer.lower() == "sgd": + return torch.optim.SGD(self.parameters(), **self.optimizer_params) + else: + raise ValueError(f"Optimizer {self.optimizer} not supported.") + elif isinstance(self.optimizer, Optimizer): + return self.optimizer + else: + raise ValueError( + "Optimizer must be either a string or " + "an instance of torch.optim.Optimizer." + ) + + def _get_scheduler( + self, optimizer: Optimizer + ) -> torch.optim.lr_scheduler._LRScheduler: + """ + Get the lr scheduler based on the specified scheduler name and params. + + Parameters + ---------- + optimizer : Optimizer + The optimizer instance. + + Returns + ------- + torch.optim.lr_scheduler._LRScheduler + The learning rate scheduler instance. + """ + if self.lr_scheduler.lower() == "reduce_lr_on_plateau": + return torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, **self.lr_scheduler_params + ) + elif self.lr_scheduler.lower() == "step_lr": + return torch.optim.lr_scheduler.StepLR( + optimizer, **self.lr_scheduler_params + ) + else: + raise ValueError(f"Scheduler {self.lr_scheduler} not supported.") + + def log_metrics( + self, y_hat: torch.Tensor, y: torch.Tensor, prefix: str = "val" + ) -> None: + """ + Log additional metrics during training, validation, or testing. + + Parameters + ---------- + y_hat : torch.Tensor + Predicted output tensor. + y : torch.Tensor + Target output tensor. + prefix : str + Prefix for the logged metrics (e.g., "train", "val", "test"). + """ + for metric in self.logging_metrics: + metric_value = metric(y_hat, y) + self.log( + f"{prefix}_{metric.__class__.__name__}", + metric_value, + on_step=False, + on_epoch=True, + prog_bar=True, + logger=True, + ) diff --git a/pytorch_forecasting/models/temporal_fusion_transformer/tft_version_two.py b/pytorch_forecasting/models/temporal_fusion_transformer/tft_version_two.py new file mode 100644 index 000000000..571d08b6e --- /dev/null +++ b/pytorch_forecasting/models/temporal_fusion_transformer/tft_version_two.py @@ -0,0 +1,229 @@ +######################################################################################## +# Disclaimer: This implementation is based on the new version of data pipeline and is +# experimental, please use with care. +######################################################################################## + +from typing import Optional, Union + +import torch +import torch.nn as nn +from torch.optim import Optimizer + +from pytorch_forecasting.models.base.base_model_refactor import BaseModel + + +class TFT(BaseModel): + def __init__( + self, + loss: nn.Module, + logging_metrics: Optional[list[nn.Module]] = None, + optimizer: Optional[Union[Optimizer, str]] = "adam", + optimizer_params: Optional[dict] = None, + lr_scheduler: Optional[str] = None, + lr_scheduler_params: Optional[dict] = None, + hidden_size: int = 64, + num_layers: int = 2, + attention_head_size: int = 4, + dropout: float = 0.1, + metadata: Optional[dict] = None, + output_size: int = 1, + ): + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params, + ) + self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"]) + + self.hidden_size = hidden_size + self.num_layers = num_layers + self.attention_head_size = attention_head_size + self.dropout = dropout + self.metadata = metadata + self.output_size = output_size + + self.max_encoder_length = self.metadata["max_encoder_length"] + self.max_prediction_length = self.metadata["max_prediction_length"] + self.encoder_cont = self.metadata["encoder_cont"] + self.encoder_cat = self.metadata["encoder_cat"] + self.encoder_input_dim = self.encoder_cont + self.encoder_cat + self.decoder_cont = self.metadata["decoder_cont"] + self.decoder_cat = self.metadata["decoder_cat"] + self.decoder_input_dim = self.decoder_cont + self.decoder_cat + self.static_cat_dim = self.metadata.get("static_categorical_features", 0) + self.static_cont_dim = self.metadata.get("static_continuous_features", 0) + self.static_input_dim = self.static_cat_dim + self.static_cont_dim + + if self.encoder_input_dim > 0: + self.encoder_var_selection = nn.Sequential( + nn.Linear(self.encoder_input_dim, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, self.encoder_input_dim), + nn.Sigmoid(), + ) + else: + self.encoder_var_selection = None + + if self.decoder_input_dim > 0: + self.decoder_var_selection = nn.Sequential( + nn.Linear(self.decoder_input_dim, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, self.decoder_input_dim), + nn.Sigmoid(), + ) + else: + self.decoder_var_selection = None + + if self.static_input_dim > 0: + self.static_context_linear = nn.Linear(self.static_input_dim, hidden_size) + else: + self.static_context_linear = None + + _lstm_encoder_input_actual_dim = self.encoder_input_dim + self.lstm_encoder = nn.LSTM( + input_size=max(1, _lstm_encoder_input_actual_dim), + hidden_size=hidden_size, + num_layers=num_layers, + dropout=dropout, + batch_first=True, + ) + + _lstm_decoder_input_actual_dim = self.decoder_input_dim + self.lstm_decoder = nn.LSTM( + input_size=max(1, _lstm_decoder_input_actual_dim), + hidden_size=hidden_size, + num_layers=num_layers, + dropout=dropout, + batch_first=True, + ) + + self.self_attention = nn.MultiheadAttention( + embed_dim=hidden_size, + num_heads=attention_head_size, + dropout=dropout, + batch_first=True, + ) + + self.pre_output = nn.Linear(hidden_size, hidden_size) + self.output_layer = nn.Linear(hidden_size, self.output_size) + + def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Forward pass of the TFT model. + + Parameters + ---------- + x : dict[str, torch.Tensor] + Dictionary containing input tensors: + - encoder_cat: Categorical encoder features + - encoder_cont: Continuous encoder features + - decoder_cat: Categorical decoder features + - decoder_cont: Continuous decoder features + - static_categorical_features: Static categorical features + - static_continuous_features: Static continuous features + + Returns + ------- + dict[str, torch.Tensor] + Dictionary containing output tensors: + - prediction: Prediction output (batch_size, prediction_length, output_size) + """ + batch_size = x["encoder_cont"].shape[0] + + encoder_cat = x.get( + "encoder_cat", + torch.zeros(batch_size, self.max_encoder_length, 0, device=self.device), + ) + encoder_cont = x.get( + "encoder_cont", + torch.zeros(batch_size, self.max_encoder_length, 0, device=self.device), + ) + decoder_cat = x.get( + "decoder_cat", + torch.zeros(batch_size, self.max_prediction_length, 0, device=self.device), + ) + decoder_cont = x.get( + "decoder_cont", + torch.zeros(batch_size, self.max_prediction_length, 0, device=self.device), + ) + + encoder_input = torch.cat([encoder_cont, encoder_cat], dim=2) + decoder_input = torch.cat([decoder_cont, decoder_cat], dim=2) + + static_context = None + if self.static_context_linear is not None: + static_cat = x.get( + "static_categorical_features", + torch.zeros(batch_size, 0, device=self.device), + ) + static_cont = x.get( + "static_continuous_features", + torch.zeros(batch_size, 0, device=self.device), + ) + + if static_cat.size(2) == 0 and static_cont.size(2) == 0: + static_context = None + elif static_cat.size(2) == 0: + static_input = static_cont.to( + dtype=self.static_context_linear.weight.dtype + ) + static_context = self.static_context_linear(static_input) + static_context = static_context.view(batch_size, self.hidden_size) + elif static_cont.size(2) == 0: + static_input = static_cat.to( + dtype=self.static_context_linear.weight.dtype + ) + static_context = self.static_context_linear(static_input) + static_context = static_context.view(batch_size, self.hidden_size) + else: + + static_input = torch.cat([static_cont, static_cat], dim=1).to( + dtype=self.static_context_linear.weight.dtype + ) + static_context = self.static_context_linear(static_input) + static_context = static_context.view(batch_size, self.hidden_size) + + encoder_weights = self.encoder_var_selection(encoder_input) + encoder_input = encoder_input * encoder_weights + + decoder_weights = self.decoder_var_selection(decoder_input) + decoder_input = decoder_input * decoder_weights + + if static_context is not None: + encoder_static_context = static_context.unsqueeze(1).expand( + -1, self.max_encoder_length, -1 + ) + decoder_static_context = static_context.unsqueeze(1).expand( + -1, self.max_prediction_length, -1 + ) + + encoder_output, (h_n, c_n) = self.lstm_encoder(encoder_input) + encoder_output = encoder_output + encoder_static_context + decoder_output, _ = self.lstm_decoder(decoder_input, (h_n, c_n)) + decoder_output = decoder_output + decoder_static_context + else: + encoder_output, (h_n, c_n) = self.lstm_encoder(encoder_input) + decoder_output, _ = self.lstm_decoder(decoder_input, (h_n, c_n)) + + sequence = torch.cat([encoder_output, decoder_output], dim=1) + + if static_context is not None: + expanded_static_context = static_context.unsqueeze(1).expand( + -1, sequence.size(1), -1 + ) + + attended_output, _ = self.self_attention( + sequence + expanded_static_context, sequence, sequence + ) + else: + attended_output, _ = self.self_attention(sequence, sequence, sequence) + + decoder_attended = attended_output[:, -self.max_prediction_length :, :] + + output = nn.functional.relu(self.pre_output(decoder_attended)) + prediction = self.output_layer(output) + + return {"prediction": prediction} diff --git a/pytorch_forecasting/models/timexer/__init__.py b/pytorch_forecasting/models/timexer/__init__.py new file mode 100644 index 000000000..8d3d51d94 --- /dev/null +++ b/pytorch_forecasting/models/timexer/__init__.py @@ -0,0 +1,29 @@ +""" +TimeXer model for forecasting time series. +""" + +from pytorch_forecasting.models.timexer._timexer import TimeXer +from pytorch_forecasting.models.timexer.sub_modules import ( + AttentionLayer, + DataEmbedding_inverted, + Encoder, + EncoderLayer, + EnEmbedding, + FlattenHead, + FullAttention, + PositionalEmbedding, + TriangularCausalMask, +) + +__all__ = [ + "TimeXer", + "TriangularCausalMask", + "FullAttention", + "AttentionLayer", + "DataEmbedding_inverted", + "PositionalEmbedding", + "FlattenHead", + "EnEmbedding", + "Encoder", + "EncoderLayer", +] diff --git a/pytorch_forecasting/models/timexer/_timexer.py b/pytorch_forecasting/models/timexer/_timexer.py new file mode 100644 index 000000000..221476c95 --- /dev/null +++ b/pytorch_forecasting/models/timexer/_timexer.py @@ -0,0 +1,289 @@ +""" +Time Series Transformer with eXogenous variables (TimeXer) +--------------------------------------------------------- +""" + +####################################################### +# Note: This is an example version to demonstrate the +# working of the TimeXer model with the exisiting v2 +# designs. The pending work includes building the D2 +# layer and base tslib model. +###################################################### + +from copy import copy +from typing import Callable, Optional, Union + +import lightning.pytorch as pl +from lightning.pytorch import LightningModule, Trainer +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.optim import Optimizer + +from pytorch_forecasting.metrics import ( + MAE, + MAPE, + MASE, + RMSE, + SMAPE, + MultiHorizonMetric, + QuantileLoss, +) +from pytorch_forecasting.models.base.base_model_refactor import BaseModel +from pytorch_forecasting.models.timexer.sub_modules import ( + AttentionLayer, + DataEmbedding_inverted, + Encoder, + EncoderLayer, + EnEmbedding, + FlattenHead, + FullAttention, +) + + +class TimeXer(BaseModel): + def __init__( + self, + context_length: int, + prediction_length: int, + loss: nn.Module, + logging_metrics: Optional[list[nn.Module]] = None, + optimizer: Optional[Union[Optimizer, str]] = "adam", + optimizer_params: Optional[dict] = None, + lr_scheduler: Optional[str] = None, + lr_scheduler_params: Optional[dict] = None, + task_name: str = "long_term_forecast", + features: str = "MS", + enc_in: int = None, + d_model: int = 512, + n_heads: int = 8, + e_layers: int = 2, + d_ff: int = 2048, + dropout: float = 0.1, + activation: Union[str, Callable] = "torch.nn.functional.relu", + patch_length: int = 24, + use_norm: bool = False, + factor: int = 5, + embed_type: str = "fixed", + freq: str = "h", + metadata: Optional[dict] = None, + target_positions: torch.LongTensor = None, + ): + """An implementation of the TimeXer model. + + TimeXer empowers the canonical transformer with the ability to reconcile + endogenous and exogenous information without any architectural modifications + and achieves consistent state-of-the-art performance across twelve real-world + forecasting benchmarks. + + TimeXer employs patch-level and variate-level representations respectively for + endogenous and exogenous variables, with an endogenous global token as a bridge + in-between. With this design, TimeXer can jointly capture intra-endogenous + temporal dependencies and exogenous-to-endogenous correlations. + + TimeXer model for time series forecasting with exogenous variables. + + """ + super().__init__( + loss=loss, + logging_metrics=logging_metrics, + optimizer=optimizer, + optimizer_params=optimizer_params or {}, + lr_scheduler=lr_scheduler, + lr_scheduler_params=lr_scheduler_params or {}, + ) + + self.context_length = context_length + self.prediction_length = prediction_length + self.task_name = task_name + self.features = features + self.d_model = d_model + self.n_heads = n_heads + self.e_layers = e_layers + self.d_ff = d_ff + self.activation = activation + self.patch_length = patch_length + self.use_norm = use_norm + self.factor = factor + self.embed_type = embed_type + self.freq = freq + self.metadata = metadata + self.n_target_vars = self.metadata["target"] + self.target_positions = target_positions + self.enc_in = self.metadata["encoder_cont"] + self.patch_num = self.context_length // self.patch_length + self.dropout = dropout + + self.n_quantiles = None + + if isinstance(loss, QuantileLoss): + self.n_quantiles = len(loss.quantiles) + + self.en_embedding = EnEmbedding( + self.n_target_vars, + self.d_model, + self.patch_length, + self.dropout, + ) + + self.ex_embedding = DataEmbedding_inverted( + self.context_length, + self.d_model, + self.embed_type, + self.freq, + self.dropout, + ) + + self.encoder = Encoder( + [ + EncoderLayer( + AttentionLayer( + FullAttention( + False, + self.factor, + attention_dropout=self.dropout, + output_attention=False, + ), + self.d_model, + self.n_heads, + ), + AttentionLayer( + FullAttention( + False, + self.factor, + attention_dropout=self.dropout, + output_attention=False, + ), + self.d_model, + self.n_heads, + ), + self.d_model, + self.d_ff, + dropout=self.dropout, + activation=self.activation, + ) + for l in range(self.e_layers) + ], + norm_layer=torch.nn.LayerNorm(self.d_model), + ) + self.head_nf = self.d_model * (self.patch_num + 1) + self.head = FlattenHead( + self.enc_in, + self.head_nf, + self.prediction_length, + head_dropout=self.dropout, + n_quantiles=self.n_quantiles, + ) + + def _forecast(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Forecast for univariate or multivariate with single target (MS) case. + + Args: + x: Dictionary containing entries for encoder_cat, encoder_cont + """ + batch_size = x["encoder_cont"].shape[0] + encoder_cont = x["encoder_cont"] + encoder_time_idx = x.get("encoder_time_idx", None) + past_target = x.get( + "target", + torch.zeros(batch_size, self.prediction_length, 0, device=self.device), + ) + + if encoder_time_idx is not None and encoder_time_idx.dim() == 2: + # change [batch_size, time_steps] to [batch_size, time_steps, features] + encoder_time_idx = encoder_time_idx.unsqueeze(-1) + + en_embed, n_vars = self.en_embedding(past_target.permute(0, 2, 1)) + ex_embed = self.ex_embedding(encoder_cont, encoder_time_idx) + + enc_out = self.encoder(en_embed, ex_embed) + enc_out = torch.reshape( + enc_out, (-1, n_vars, enc_out.shape[-2], enc_out.shape[-1]) + ) + + enc_out = enc_out.permute(0, 1, 3, 2) + + dec_out = self.head(enc_out) + if self.n_quantiles is not None: + dec_out = dec_out.permute(0, 2, 1, 3) + else: + dec_out = dec_out.permute(0, 2, 1) + + return dec_out + + def _forecast_multi(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Forecast for multivariate with multiple targets (M) case. + + Args: + x: Dictionary containing entries for encoder_cat, encoder_cont + Returns: + Dictionary with predictions + """ + + batch_size = x["encoder_cont"].shape[0] + encoder_cont = x.get( + "encoder_cont", + torch.zeros(batch_size, self.prediction_length, device=self.device), + ) + encoder_time_idx = x.get("encoder_time_idx", None) + encoder_targets = x.get( + "target", + torch.zeros(batch_size, self.prediction_length, device=self.device), + ) + en_embed, n_vars = self.en_embedding(encoder_targets.permute(0, 2, 1)) + ex_embed = self.ex_embedding(encoder_cont, encoder_time_idx) + + # batch_size x sequence_length x d_model + enc_out = self.encoder(en_embed, ex_embed) + + enc_out = torch.reshape( + enc_out, (-1, n_vars, enc_out.shape[-2], enc_out.shape[-1]) + ) # batch_size x n_vars x sequence_length x d_model + + enc_out = enc_out.permute(0, 1, 3, 2) + + dec_out = self.head(enc_out) + if self.n_quantiles is not None: + dec_out = dec_out.permute(0, 2, 1, 3) + else: + dec_out = dec_out.permute(0, 2, 1) + + return dec_out + + def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Forward pass of the model. + + Args: + x: Dictionary containing model inputs + + Returns: + Dictionary with model outputs + """ + if ( + self.task_name == "long_term_forecast" + or self.task_name == "short_term_forecast" + ): # noqa: E501 + if self.features == "M": + out = self._forecast_multi(x) + else: + out = self._forecast(x) + prediction = out[:, : self.prediction_length, :] + + # note: prediction.size(2) is the number of target variables i.e n_targets + target_indices = range(prediction.size(2)) + + if self.n_quantiles is not None: + prediction = [prediction[..., i, :] for i in target_indices] + else: + + if len(target_indices) == 1: + prediction = prediction[..., 0] + else: + prediction = [prediction[..., i] for i in target_indices] + return {"prediction": prediction} + else: + return None diff --git a/pytorch_forecasting/models/timexer/sub_modules.py b/pytorch_forecasting/models/timexer/sub_modules.py new file mode 100644 index 000000000..c13b9fc61 --- /dev/null +++ b/pytorch_forecasting/models/timexer/sub_modules.py @@ -0,0 +1,251 @@ +""" +Implementation of `nn.Modules` for TimeXer model. +""" + +import math +from math import sqrt + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class TriangularCausalMask: + def __init__(self, B, L, device="cpu"): + mask_shape = [B, 1, L, L] + with torch.no_grad(): + self._mask = torch.triu( + torch.ones(mask_shape, dtype=torch.bool), diagonal=1 + ).to(device) + + @property + def mask(self): + return self._mask + + +class FullAttention(nn.Module): + def __init__( + self, + mask_flag=True, + factor=5, + scale=None, + attention_dropout=0.1, + output_attention=False, + ): + super().__init__() + self.scale = scale + self.mask_flag = mask_flag + self.output_attention = output_attention + self.dropout = nn.Dropout(attention_dropout) + + def forward(self, queries, keys, values, attn_mask, tau=None, delta=None): + B, L, H, E = queries.shape + _, S, _, D = values.shape + scale = self.scale or 1.0 / sqrt(E) + + scores = torch.einsum("blhe,bshe->bhls", queries, keys) + + if self.mask_flag: + if attn_mask is None: + attn_mask = TriangularCausalMask(B, L, device=queries.device) + scores.masked_fill_(attn_mask.mask, -np.abs) + A = self.dropout(torch.softmax(scale * scores, dim=-1)) + V = torch.einsum("bhls,bshd->blhd", A, values) + + if self.output_attention: + return V.contiguous(), A + else: + return V.contiguous(), None + + +class AttentionLayer(nn.Module): + def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None): + super().__init__() + + d_keys = d_keys or (d_model // n_heads) + d_values = d_values or (d_model // n_heads) + + self.inner_attention = attention + self.query_projection = nn.Linear(d_model, d_keys * n_heads) + self.key_projection = nn.Linear(d_model, d_keys * n_heads) + self.value_projection = nn.Linear(d_model, d_values * n_heads) + self.out_projection = nn.Linear(d_values * n_heads, d_model) + self.n_heads = n_heads + + def forward(self, queries, keys, values, attn_mask, tau=None, delta=None): + B, L, _ = queries.shape + _, S, _ = keys.shape + H = self.n_heads + + queries = self.query_projection(queries).view(B, L, H, -1) + keys = self.key_projection(keys).view(B, S, H, -1) + values = self.value_projection(values).view(B, S, H, -1) + + out, attn = self.inner_attention( + queries, keys, values, attn_mask, tau=tau, delta=delta + ) + out = out.view(B, L, -1) + + return self.out_projection(out), attn + + +class DataEmbedding_inverted(nn.Module): + def __init__(self, c_in, d_model, embed_type="fixed", freq="h", dropout=0.1): + super().__init__() + self.value_embedding = nn.Linear(c_in, d_model) + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x, x_mark): + x = x.permute(0, 2, 1) + # x: [Batch Variate Time] + if x_mark is None: + x = self.value_embedding(x) + else: + x = self.value_embedding(torch.cat([x, x_mark.permute(0, 2, 1)], 1)) + # x: [Batch Variate d_model] + return self.dropout(x) + + +class PositionalEmbedding(nn.Module): + def __init__(self, d_model, max_len=5000): + super().__init__() + # Compute the positional encodings once in log space. + pe = torch.zeros(max_len, d_model).float() + pe.require_grad = False + + position = torch.arange(0, max_len).float().unsqueeze(1) + div_term = ( + torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model) + ).exp() + + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + + pe = pe.unsqueeze(0) + self.register_buffer("pe", pe) + + def forward(self, x): + return self.pe[:, : x.size(1)] + + +class FlattenHead(nn.Module): + def __init__(self, n_vars, nf, target_window, head_dropout=0, n_quantiles=None): + super().__init__() + self.n_vars = n_vars + self.flatten = nn.Flatten(start_dim=-2) + self.linear = nn.Linear(nf, target_window) + self.n_quantiles = n_quantiles + + if self.n_quantiles is not None: + self.linear = nn.Linear(nf, target_window * n_quantiles) + else: + self.linear = nn.Linear(nf, target_window) + self.dropout = nn.Dropout(head_dropout) + + def forward(self, x): + x = self.flatten(x) + x = self.linear(x) + x = self.dropout(x) + + if self.n_quantiles is not None: + batch_size, n_vars = x.shape[0], x.shape[1] + x = x.reshape(batch_size, n_vars, -1, self.n_quantiles) + return x + + +class EnEmbedding(nn.Module): + def __init__(self, n_vars, d_model, patch_len, dropout): + super().__init__() + + self.patch_len = patch_len + + self.value_embedding = nn.Linear(patch_len, d_model, bias=False) + self.glb_token = nn.Parameter(torch.randn(1, n_vars, 1, d_model)) + self.position_embedding = PositionalEmbedding(d_model) + + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + n_vars = x.shape[1] + glb = self.glb_token.repeat((x.shape[0], 1, 1, 1)) + + x = x.unfold(dimension=-1, size=self.patch_len, step=self.patch_len) + x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3])) + # Input encoding + x = self.value_embedding(x) + self.position_embedding(x) + x = torch.reshape(x, (-1, n_vars, x.shape[-2], x.shape[-1])) + x = torch.cat([x, glb], dim=2) + x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3])) + return self.dropout(x), n_vars + + +class Encoder(nn.Module): + def __init__(self, layers, norm_layer=None, projection=None): + super().__init__() + self.layers = nn.ModuleList(layers) + self.norm = norm_layer + self.projection = projection + + def forward(self, x, cross, x_mask=None, cross_mask=None, tau=None, delta=None): + for layer in self.layers: + x = layer( + x, cross, x_mask=x_mask, cross_mask=cross_mask, tau=tau, delta=delta + ) + + if self.norm is not None: + x = self.norm(x) + + if self.projection is not None: + x = self.projection(x) + return x + + +class EncoderLayer(nn.Module): + def __init__( + self, + self_attention, + cross_attention, + d_model, + d_ff=None, + dropout=0.1, + activation="relu", + ): + super().__init__() + d_ff = d_ff or 4 * d_model + self.self_attention = self_attention + self.cross_attention = cross_attention + self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1) + self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + self.activation = F.relu if activation == "relu" else F.gelu + + def forward(self, x, cross, x_mask=None, cross_mask=None, tau=None, delta=None): + B, L, D = cross.shape + x = x + self.dropout( + self.self_attention(x, x, x, attn_mask=x_mask, tau=tau, delta=None)[0] + ) + x = self.norm1(x) + + x_glb_ori = x[:, -1, :].unsqueeze(1) + x_glb = torch.reshape(x_glb_ori, (B, -1, D)) + x_glb_attn = self.dropout( + self.cross_attention( + x_glb, cross, cross, attn_mask=cross_mask, tau=tau, delta=delta + )[0] + ) + x_glb_attn = torch.reshape( + x_glb_attn, (x_glb_attn.shape[0] * x_glb_attn.shape[1], x_glb_attn.shape[2]) + ).unsqueeze(1) + x_glb = x_glb_ori + x_glb_attn + x_glb = self.norm2(x_glb) + + y = x = torch.cat([x[:, :-1, :], x_glb], dim=1) + + y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1)))) + y = self.dropout(self.conv2(y).transpose(-1, 1)) + + return self.norm3(x + y)