freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

2.0.0 initial

Browse Source
This commit is contained in:
siriuz42
2025-09-12 00:18:08 +00:00
parent d70708d42a
commit 7d8f3d971d
52 changed files with 1882 additions and 394 deletions
Download Patch File Download Diff File Expand all files Collapse all files
v1/LICENSE
+202
View File
@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
v1/README.md
+313
View File
@@ -0,0 +1,313 @@
# TimesFM
TimesFM (Time Series Foundation Model) is a pretrained time-series foundation model developed by Google
Research for time-series forecasting.
* Paper: [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688), to appear in ICML 2024.
* [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/)
* [Hugging Face release](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6)
This repo contains the code to load public TimesFM checkpoints and run model
inference. Please visit our
[Hugging Face release](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6)
to download model checkpoints.
This is not an officially supported Google product.
We recommend at least 32GB RAM to load TimesFM dependencies.
**Need help?** See [TROUBLESHOOTING.md](TROUBLESHOOTING.md) for common installation and usage issues.
## Update - Dec. 30, 2024
- We are launching a 500m checkpoint as a part of TimesFM-2.0 release. This new checkpoint can be upto 25% better than v1.0 on leading benchmarks and also has a 4 times longer max. context length.
- Launched [finetuning support](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning.ipynb) that lets you finetune the weights of the pretrained TimesFM model on your own data.
- Launched [~zero-shot covariate support](https://github.com/google-research/timesfm/blob/master/notebooks/covariates.ipynb) with external regressors. More details [here](https://github.com/google-research/timesfm?tab=readme-ov-file#covariates-support).
## Update - Feb. 17, 2024
- We are providing the option for [finetuning using Pytorch](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning_torch.ipynb), which mimics the previously added functionality from [finetuning support](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning.ipynb).
- We are also providing the Multi-GPU finetuining with Pytorch. We currently support DDP multi-gpu finetuning, other variants of multi-gpu training (pipeline parallelism/model parallelism) might be added later. In order to use it, follow the steps in [finetuning example](https://github.com/google-research/timesfm/blob/master/finetuning/finetuning_example.py) .
## Checkpoint timesfm-1.0-200m (-pytorch)
timesfm-1.0-200m is our first open model checkpoint:
- It performs univariate time series forecasting for context lengths up to 512 timepoints and any horizon lengths, with an optional frequency indicator.
- It focuses on point forecasts, and does not support probabilistic forecasts. We experimentally offer quantile heads but they have not been calibrated after pretraining.
## Checkpoint timesfm-2.0-500m (-jax/-pytorch)
timesfm-2.0-500m is our second open model checkpoint:
- It performs univariate time series forecasting for context lengths up to 2048 timepoints and any horizon lengths, with an optional frequency indicator.
- It focuses on point forecasts. We experimentally offer 10 quantile heads but they have not been calibrated after pretraining.
- This new checkpoint can be upto 25% better than v1.0 on leading benchmarks and also has a 4 times longer max. context length.
## Benchmarking
TimesFM 2.0 has been added to [GIFT-Eval](https://huggingface.co/spaces/Salesforce/GIFT-Eval) which is one of the most comprehensive time-series bechmarks available. It takes the top spot in terms of aggregated MASE and CRPS, where it is 6\% better than the next best model in terms of aggregated MASE.
## Installation
### Local installation using poetry
We will be using `pyenv` and `poetry`. In order to set these things up please follow the instructions [here](https://substack.com/home/post/p-148747960?r=28a5lx&utm_campaign=post&utm_medium=web). Note that the PAX (or JAX) version needs to run on python 3.10.x and the PyTorch version can run on >=3.11.x. Therefore make sure you have two versions of python installed:
```
pyenv install 3.10
pyenv install 3.11
pyenv versions # to list the versions available (lets assume the versions are 3.10.15 and 3.11.10)
```
### For PAX version installation do the following.
```
pyenv local 3.10.15
poetry env use 3.10.15
poetry lock
poetry install -E pax
```
After than you can run the timesfm under `poetry shell` or do `poetry run python3 ...`.
### For PyTorch version installation do the following.
```
pyenv local 3.11.10
poetry env use 3.11.10
poetry lock
poetry install -E torch
```
After than you can run the timesfm under `poetry shell` or do `poetry run python3 ...`.
**Additional Note**:
If you plan to use the **`forecast_with_covariates`** function (which requires external regressors),
you need to install **JAX** and **jaxlib**. If you installed the base version of TimesFM (`torch`), you must manually install the dependencies for **`forecast_with_covariates`** support:
```
pip install jax jaxlib
```
**Why is this needed?**
The `forecast_with_covariates` method relies on the `xreg_lib` module, which depends on JAX and jaxlib. If these packages are not installed,
calling `forecast_with_covariates` will raise an error. However, due to a lazy import mechanism, `xreg_lib` (and hence JAX/jaxlib) is not needed for standard `forecast` calls.
### Notes
1. Running the provided benchmarks would require additional dependencies. Please see the `experiments` folder.
2. The dependency `lingvo` does not support ARM architectures, and the code is not working for machines with Apple silicon. We are aware of this issue and are working on a solution. Stay tuned.
### Install from PyPI (and publish)
On python 3.11 you can install the torch version using:
```pip install timesfm[torch]```
On python 3.10 you can install the pax version using:
```pip install timesfm[pax]```
## Usage
### Initialize the model and load a checkpoint.
Then the base class can be loaded as,
```python
import timesfm
# Loading the timesfm-2.0 checkpoint:
# For PAX
tfm = timesfm.TimesFm(
hparams=timesfm.TimesFmHparams(
backend="gpu",
per_core_batch_size=32,
horizon_len=128,
num_layers=50,
context_len=2048,
use_positional_embedding=False,
),
checkpoint=timesfm.TimesFmCheckpoint(
huggingface_repo_id="google/timesfm-2.0-500m-jax"),
)
# For Torch
tfm = timesfm.TimesFm(
hparams=timesfm.TimesFmHparams(
backend="gpu",
per_core_batch_size=32,
horizon_len=128,
num_layers=50,
use_positional_embedding=False,
context_len=2048,
),
checkpoint=timesfm.TimesFmCheckpoint(
huggingface_repo_id="google/timesfm-2.0-500m-pytorch"),
)
# Loading the timesfm-1.0 checkpoint:
# For PAX
tfm = timesfm.TimesFm(
hparams=timesfm.TimesFmHparams(
backend="gpu",
per_core_batch_size=32,
horizon_len=128,
),
checkpoint=timesfm.TimesFmCheckpoint(
huggingface_repo_id="google/timesfm-1.0-200m"),
)
# For Torch
tfm = timesfm.TimesFm(
hparams=timesfm.TimesFmHparams(
backend="gpu",
per_core_batch_size=32,
horizon_len=128,
),
checkpoint=timesfm.TimesFmCheckpoint(
huggingface_repo_id="google/timesfm-1.0-200m-pytorch"),
)
```
Note some of the parameters are fixed to load the 200m and 500m models
1. The `context_len` in `hparams` here can be set as the max context length **of the model** (a maximum of 2048 for 2.0 models and 512 for 1.0 models). **It needs to be a multiplier of `input_patch_len`, i.e. a multiplier of 32.** You can provide a shorter series to the `tfm.forecast()` function and the model will handle it. The input time series can have **any context length**. Padding / truncation will be handled by the inference code if needed.
2. The horizon length can be set to anything. We recommend setting it to the largest horizon length you would need in the forecasting tasks for your application. We generally recommend horizon length <= context length but it is not a requirement in the function call.
3. `backend` is one of "cpu", "gpu", case sensitive.
### Perform inference
We provide APIs to forecast from either array inputs or `pandas` dataframe. Both forecast methods expect (1) the input time series contexts, (2) along with their frequencies. Please look at the documentation of the functions `tfm.forecast()` and `tfm.forecast_on_df()` for detailed instructions.
In particular regarding the frequency, TimesFM expects a categorical indicator valued in {0, 1, 2}:
- **0** (default): high frequency, long horizon time series. We recommend using this for time series up to daily granularity.
- **1**: medium frequency time series. We recommend using this for weekly and monthly data.
- **2**: low frequency, short horizon time series. We recommend using this for anything beyond monthly, e.g. quarterly or yearly.
This categorical value should be directly provided with the array inputs. For dataframe inputs, we convert the conventional letter coding of frequencies to our expected categories, that
- **0**: T, MIN, H, D, B, U
- **1**: W, M
- **2**: Q, Y
Notice you do **NOT** have to strictly follow our recommendation here. Although this is our setup during model training and we expect it to offer the best forecast result, you can also view the frequency input as a free parameter and modify it per your specific use case.
Examples:
Array inputs, with the frequencies set to low, medium and high respectively.
```python
import numpy as np
forecast_input = [
np.sin(np.linspace(0, 20, 100)),
np.sin(np.linspace(0, 20, 200)),
np.sin(np.linspace(0, 20, 400)),
]
frequency_input = [0, 1, 2]
point_forecast, experimental_quantile_forecast = tfm.forecast(
forecast_input,
freq=frequency_input,
)
```
`pandas` dataframe, with the frequency set to "M" monthly.
```python
import pandas as pd
# e.g. input_df is
# unique_id ds y
# 0 T1 1975-12-31 697458.0
# 1 T1 1976-01-31 1187650.0
# 2 T1 1976-02-29 1069690.0
# 3 T1 1976-03-31 1078430.0
# 4 T1 1976-04-30 1059910.0
# ... ... ... ...
# 8175 T99 1986-01-31 602.0
# 8176 T99 1986-02-28 684.0
# 8177 T99 1986-03-31 818.0
# 8178 T99 1986-04-30 836.0
# 8179 T99 1986-05-31 878.0
forecast_df = tfm.forecast_on_df(
inputs=input_df,
freq="M", # monthly
value_name="y",
num_jobs=-1,
)
```
## Covariates Support
We now have an external regressors library on top of TimesFM that can support static covariates as well as dynamic covariates available in the future. We have an usage example in [notebooks/covariates.ipynb](https://github.com/google-research/timesfm/blob/master/notebooks/covariates.ipynb).
If you plan to use the **`forecast_with_covariates`** on timesfm `torch` version, you need to install **JAX** and **jaxlib**.
You must manually install the dependencies for **`forecast_with_covariates`** support:
```
pip install jax jaxlib
```
Let's take a toy example of forecasting sales for a grocery store:
**Task:** Given the observed the daily sales of this week (7 days), forecast the daily sales of next week (7 days).
```
Product: ice cream
Daily_sales: [30, 30, 4, 5, 7, 8, 10]
Category: food
Base_price: 1.99
Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]
Has_promotion: [Yes, Yes, No, No, No, Yes, Yes, No, No, No, No, No, No, No]
Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]
```
```
Product: sunscreen
Daily_sales: [5, 7, 12, 13, 5, 6, 10]
Category: skin product
Base_price: 29.99
Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]
Has_promotion: [No, No, Yes, Yes, No, No, No, Yes, Yes, Yes, Yes, Yes, Yes, Yes]
Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]
```
In this example, besides the `Daily_sales`, we also have covariates `Category`, `Base_price`, `Weekday`, `Has_promotion`, `Daily_temperature`. Let's introduce some concepts:
**Static covariates** are covariates for each time series.
- In our example, `Category` is a **static categorical covariate**,
- `Base_price` is a **static numerical covariates**.
**Dynamic covariates** are covaraites for each time stamps.
- Date / time related features can be usually treated as dynamic covariates.
- In our example, `Weekday` and `Has_promotion` are **dynamic categorical covariates**.
- `Daily_temperate` is a **dynamic numerical covariate**.
**Notice:** Here we make it mandatory that the dynamic covariates need to cover both the forecasting context and horizon. For example, all dynamic covariates in the example have 14 values: the first 7 correspond to the observed 7 days, and the last 7 correspond to the next 7 days.
We can now provide the past data of the two products along with static and dynamic covariates as a batch input to TimesFM and produce forecasts that take into the account the covariates. To learn more, check out the example in [notebooks/covariates.ipynb](https://github.com/google-research/timesfm/blob/master/notebooks/covariates.ipynb).
## Finetuning
We have provided an example of finetuning the model on a new dataset in [notebooks/finetuning.ipynb](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning.ipynb).
## Contribution Style guide
If you would like to submit a PR please make sure that you use our formatting style. We use [yapf](https://github.com/google/yapf) for formatting with the following options,
```
[style]
based_on_style = google
# Add your custom style rules here
indent_width = 2
spaces_before_comment = 2
```
Please run `yapf --in-place --recursive <filename>` on all affected files.
v1/TROUBLESHOOTING.md
+143
View File
@@ -0,0 +1,143 @@
# Troubleshooting
This document provides solutions to common issues encountered when using TimesFM.
## Installation Issues
### ARM/Apple Silicon Compatibility
**Problem:** `lingvo` dependency fails on Apple Silicon (M1/M2/M3) machines.
```
ERROR: Could not build wheels for lingvo
```
**Solution:** This is a known issue. The `lingvo` dependency doesn't support ARM architectures. We recommend:
- Use x86_64 emulation via Rosetta 2: `arch -x86_64 pip install timesfm[pax]`
- Use the PyTorch version instead, which has better ARM support: `pip install timesfm[torch]`
- Use Docker with x86_64 emulation for consistent environments
### Memory Issues During Installation
**Problem:** Installation fails with memory errors.
```
Killed (signal 9)
```
**Solution:**
- Ensure at least 32GB RAM is available
- Close other applications during installation
- Use `pip install --no-cache-dir timesfm[torch]` to reduce memory usage
- Install in a clean virtual environment
### JAX/PyTorch Version Conflicts
**Problem:** Conflicting JAX and PyTorch installations.
```
ImportError: cannot import name 'jax' from 'jax'
```
**Solution:**
- For PyTorch-only usage: `pip install timesfm[torch]`
- For covariates with PyTorch: `pip install timesfm[torch] && pip install jax jaxlib`
- For PAX version: `pip install timesfm[pax]`
## Runtime Errors
### Model Loading Issues
**Problem:** Checkpoint download fails or is corrupted.
```
HfFileNotFoundError: 404 Client Error
```
**Solution:**
- Check internet connectivity
- Verify Hugging Face Hub access: `huggingface-cli login`
- Clear cache: `rm -rf ~/.cache/huggingface/`
- Use explicit checkpoint paths if needed
### CUDA/GPU Issues
**Problem:** GPU not detected or CUDA errors.
```
RuntimeError: CUDA out of memory
```
**Solutions:**
- Reduce `per_core_batch_size` (try 16, 8, or 4)
- Reduce `context_len` to minimum needed
- Use `backend="cpu"` for testing
- Check GPU memory: `nvidia-smi`
### Context Length Errors
**Problem:** Input series longer than model capacity.
```
ValueError: context_len must be <= 512 for v1.0 models
```
**Solutions:**
- Use TimesFM-2.0 for longer contexts (up to 2048)
- Ensure `context_len` is multiple of 32
- Truncate input series if necessary
- Set appropriate `context_len` in model initialization
## Data Issues
### Frequency Mapping Problems
**Problem:** Unexpected forecasting results with wrong frequency.
```
Warning: Frequency 'D' mapped to category 0
```
**Solutions:**
- Verify frequency mapping: D→0 (high), W/M→1 (medium), Q/Y→2 (low)
- Override automatic mapping by specifying frequency manually
- Check data granularity matches chosen frequency category
### Missing Values in Time Series
**Problem:** NaN or missing values in input data.
```
ValueError: Input contains NaN values
```
**Solutions:**
- Pre-process data to handle missing values (forward fill, interpolation)
- Ensure continuous time series without gaps
- Remove or impute missing values before forecasting
### Covariate Dimension Mismatches
**Problem:** Covariate lengths don't match forecast horizon.
```
ValueError: Dynamic covariates must cover context + horizon
```
**Solutions:**
- Ensure dynamic covariates have length = context + horizon
- Check static vs dynamic covariate classification
- Verify covariate data alignment with time series
## Performance Issues
### Slow Inference
**Problem:** Forecasting takes unexpectedly long.
**Solutions:**
- Use GPU backend: `backend="gpu"`
- Optimize batch size: increase `per_core_batch_size`
- Use appropriate model size for your use case
- Profile with smaller data first
### Memory Usage
**Problem:** High memory consumption during inference.
**Solutions:**
- Reduce batch size: `per_core_batch_size=1`
- Process data in chunks
- Use smaller context length when possible
- Monitor memory with `htop` or `nvidia-smi`
## Common Error Messages
### `ModuleNotFoundError: No module named 'xreg_lib'`
**Cause:** Missing JAX dependencies for covariates functionality.
**Solution:** `pip install jax jaxlib`
### `ValueError: horizon_len must be positive`
**Cause:** Invalid horizon length specified.
**Solution:** Set `horizon_len > 0` in model initialization.
### `RuntimeError: Expected input batch_size (X) to be divisible by batch_size (Y)`
**Cause:** Batch size mismatch.
**Solution:** Adjust `per_core_batch_size` or input data batching.
## Getting Help
If you encounter issues not covered here:
1. Check the [GitHub Issues](https://github.com/google-research/timesfm/issues)
2. Review the [notebooks/](notebooks/) for working examples
3. Verify your installation follows the exact steps in the Installation section
4. Test with the provided example data before using your own datasets
v1/docs/contributing.md
+32
View File
@@ -0,0 +1,32 @@
# How to Contribute
We would love to accept your patches and contributions to this project.
## Before you begin
### Sign our Contributor License Agreement
Contributions to this project must be accompanied by a
[Contributor License Agreement](https://cla.developers.google.com/about) (CLA).
You (or your employer) retain the copyright to your contribution; this simply
gives us permission to use and redistribute your contributions as part of the
project.
If you or your current employer have already signed the Google CLA (even if it
was for a different project), you probably don't need to do it again.
Visit <https://cla.developers.google.com/> to see your current agreements or to
sign a new one.
### Review our Community Guidelines
This project follows [Google's Open Source Community
Guidelines](https://opensource.google/conduct/).
## Contribution process
### Code Reviews
All submissions, including submissions by project members, require review. We
use [GitHub pull requests](https://docs.github.com/articles/about-pull-requests)
for this purpose.
v1/experiments/baselines/__init__.py
+13
View File
@@ -0,0 +1,13 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
v1/experiments/baselines/timegpt_pipeline.py
+259
View File
@@ -0,0 +1,259 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from time import time
from typing import List, Optional, Tuple
from dotenv import load_dotenv
from gluonts.time_feature.seasonality import get_seasonality as _get_seasonality
from nixtla import NixtlaClient
import pandas as pd
from tqdm import tqdm
from utilsforecast.processing import (
backtest_splits,
drop_index_if_pandas,
join,
maybe_compute_sort_indices,
take_rows,
vertical_concat,
)
def get_seasonality(freq: str) -> int:
return _get_seasonality(freq, seasonalities={"D": 7})
def maybe_convert_col_to_datetime(
df: pd.DataFrame, col_name: str
) -> pd.DataFrame:
if not pd.api.types.is_datetime64_any_dtype(df[col_name]):
df = df.copy()
df[col_name] = pd.to_datetime(df[col_name])
return df
def zero_pad_time_series(df, freq, min_length=36):
"""If time_series length is less than min_length, front pad it with zeros."""
# 1. Calculate required padding for each unique_id
value_counts = df["unique_id"].value_counts()
to_pad = value_counts[value_counts < min_length].index
# 2. Create a new DataFrame to hold padded data
padded_data = []
for unique_id in to_pad:
# 2a. Filter data for the specific unique_id
subset = df[df["unique_id"] == unique_id]
if len(subset) > min_length:
padded_data.append(subset)
else:
# 2b. Determine earliest date and calculate padding dates
start_date = subset["ds"].min()
padding_dates = pd.date_range(
end=start_date,
periods=min_length - len(subset) + 1,
freq=freq, # 'MS' for month start
)[
:-1
] # Exclude the start_date itself
# 2c. Create padding data
padding_df = pd.DataFrame(
{"ds": padding_dates, "unique_id": unique_id, "y": 0} # Zero padding
)
# 2d. Combine original and padding data, and append to the list
padded_data.append(pd.concat([padding_df, subset]).sort_values("ds"))
# 3. Combine all padded data and original data (unchanged)
result_df = pd.concat(padded_data + [df[~df["unique_id"].isin(to_pad)]])
return result_df
class Forecaster:
"""Borrowed from
https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
"""
def forecast(
self,
df: pd.DataFrame,
h: int,
freq: str,
) -> pd.DataFrame:
raise NotImplementedError
def cross_validation(
self,
df: pd.DataFrame,
h: int,
freq: str,
n_windows: int = 1,
step_size: int | None = None,
) -> pd.DataFrame:
df = maybe_convert_col_to_datetime(df, "ds")
# mlforecast cv code
results = []
sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds")
if sort_idxs is not None:
df = take_rows(df, sort_idxs)
splits = backtest_splits(
df,
n_windows=n_windows,
h=h,
id_col="unique_id",
time_col="ds",
freq=pd.tseries.frequencies.to_offset(freq),
step_size=h if step_size is None else step_size,
)
for _, (cutoffs, train, valid) in tqdm(enumerate(splits)):
if len(valid.columns) > 3:
raise NotImplementedError(
"Cross validation with exogenous variables is not yet supported."
)
y_pred = self.forecast(
df=train,
h=h,
freq=freq,
)
y_pred = join(y_pred, cutoffs, on="unique_id", how="left")
result = join(
valid[["unique_id", "ds", "y"]],
y_pred,
on=["unique_id", "ds"],
)
if result.shape[0] < valid.shape[0]:
raise ValueError(
"Cross validation result produced less results than expected."
" Please verify that the frequency parameter (freq) matches your"
" series' and that there aren't any missing periods."
)
results.append(result)
out = vertical_concat(results)
out = drop_index_if_pandas(out)
first_out_cols = ["unique_id", "ds", "cutoff", "y"]
remaining_cols = [c for c in out.columns if c not in first_out_cols]
fcst_cv_df = out[first_out_cols + remaining_cols]
return fcst_cv_df
class TimeGPT(Forecaster):
"""Borrowed from
https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
We modify the class to take care of edge cases.
"""
def __init__(
self,
api_key: str | None = None,
base_url: Optional[str] = None,
max_retries: int = 1,
model: str = "timegpt-1",
alias: str = "TimeGPT",
):
self.api_key = api_key
self.base_url = base_url
self.max_retries = max_retries
self.model = model
self.alias = alias
def _get_client(self) -> NixtlaClient:
if self.api_key is None:
api_key = os.environ["NIXTLA_API_KEY"]
else:
api_key = self.api_key
return NixtlaClient(
api_key=api_key,
base_url=self.base_url,
max_retries=self.max_retries,
)
def forecast(
self,
df: pd.DataFrame,
h: int,
freq: str,
level: List = [90.0],
chunk_size: Optional[int] = None,
) -> pd.DataFrame:
client = self._get_client()
fcst_df = None
if chunk_size is None:
fcst_df = client.forecast(
df=df,
h=h,
freq=freq,
level=level,
model=self.model,
)
else:
all_unique_ids = df["unique_id"].unique()
all_fcst_df = []
for i in range(0, len(all_unique_ids), chunk_size):
chunk_ids = all_unique_ids[i : i + chunk_size]
chunk_df = df[df["unique_id"].isin(chunk_ids)]
fct_chunk_df = client.forecast(
df=chunk_df,
h=h,
freq=freq,
level=level,
)
all_fcst_df.append(fct_chunk_df)
fcst_df = pd.concat(all_fcst_df)
fcst_df["ds"] = pd.to_datetime(fcst_df["ds"])
replace_dict = {}
for col in fcst_df.columns:
if col.startswith("TimeGPT"):
replace_dict[col] = col.replace("TimeGPT", self.alias)
fcst_df = fcst_df.rename(columns=replace_dict)
return fcst_df
def run_timegpt(
train_df: pd.DataFrame,
horizon: int,
freq: str,
seasonality: int,
level: List[int],
dataset: str,
model: str = "timegpt-1",
) -> Tuple[pd.DataFrame, float, str]:
os.environ["NIXTLA_ID_AS_COL"] = "true"
model = TimeGPT(model="timegpt-1", alias=model)
padded_train_df = zero_pad_time_series(train_df, freq)
init_time = time()
# For these datasets the API fails if we do not chunk.
if dataset in ["m5", "m4_quarterly"]:
chunk_size = 5000
else:
chunk_size = None
fcsts_df = model.forecast(
df=padded_train_df,
h=horizon,
level=level,
freq=freq,
chunk_size=chunk_size,
)
total_time = time() - init_time
# In case levels are not returned we replace the levels with the mean predictions.
# Note that this does not affect the results table as we only compare on point
# forecastign metrics.
for lvl in level:
if f"{model.alias}-lo-{lvl}" not in fcsts_df.columns:
fcsts_df[f"{model.alias}-lo-{lvl}"] = fcsts_df[model.alias]
if f"{model.alias}-hi-{lvl}" not in fcsts_df.columns:
fcsts_df[f"{model.alias}-hi-{lvl}"] = fcsts_df[model.alias]
return fcsts_df, total_time, model.alias
v1/experiments/extended_benchmarks/README.md
+35
View File
@@ -0,0 +1,35 @@
# Extended Benchmarks
The benchmark setting has been borrowed from Nixtla's original [benchmarking](https://github.com/AzulGarza/nixtla/tree/main/experiments/amazon-chronos) of time-series foundation models against a strong statistical ensemble. Later more datasets were added by the Chronos team in this [pull request](https://github.com/shchur/nixtla/tree/chronos-full-eval/experiments/amazon-chronos). We compare on all the datasets in this extended benchmarks.
## Running TimesFM on the benchmark
We need to add the following packages for running these benchmarks. Follow the installation instructions till before `poetry lock`. Then,
```
poetry add git+https://github.com/awslabs/gluon-ts.git
poetry lock
poetry install --only <pax or pytorch>
```
To run the timesfm on the benchmark do:
```
poetry run python3 -m experiments.extended_benchmarks.run_timesfm --model_path=google/timesfm-1.0-200m(-pytorch) --backend="gpu"
```
Note: In the current version of TimesFM we focus on point forecasts and therefore the mase, smape have been calculated using the quantile head corresponding to the median i.e 0.5 quantile. We do offer 10 quantile heads but they have not been calibrated after pretraining. We recommend using them with caution or calibrate/conformalize them on a hold out for your applications. More to follow on later versions.
## Benchmark Results for TimesFM-1.0
![Benchmark Results Table](./tfm_extended_new.png)
__Update:__ We have added TimeGPT-1 to the benchmark results. We had to remove the Dominick dataset as we were not able to run TimeGPT-1 on this benchmark. Note that the previous results including Dominick remain available at `./tfm_results.png`. In order to reproduce the results for TimeGPT-1, please run `run_timegpt.py`.
_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT. The TimesFM latency numbers are from the PAX version.
We can see that TimesFM performs the best in terms of both mase and smape. More importantly it is much faster than the other methods, in particular it is more than 600x faster than StatisticalEnsemble and 80x faster than Chronos (Large).
Note: This benchmark only compares on `one` small horizon window for long horizon datasets like ETT hourly and 15 minutes. More in depth comparison on longer horizon rolling validation tasks are presented in our long horizon benchmarks.
v1/experiments/extended_benchmarks/run_timegpt.py
+108
View File
@@ -0,0 +1,108 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Evaluation script for timegpt."""
import os
import sys
import time
from absl import flags
import numpy as np
import pandas as pd
from ..baselines.timegpt_pipeline import run_timegpt
from .utils import ExperimentHandler
dataset_names = [
"m1_monthly",
"m1_quarterly",
"m1_yearly",
"m3_monthly",
"m3_other",
"m3_quarterly",
"m3_yearly",
"m4_quarterly",
"m4_yearly",
"tourism_monthly",
"tourism_quarterly",
"tourism_yearly",
"nn5_daily_without_missing",
"m5",
"nn5_weekly",
"traffic",
"weather",
"australian_electricity_demand",
"car_parts_without_missing",
"cif_2016",
"covid_deaths",
"ercot",
"ett_small_15min",
"ett_small_1h",
"exchange_rate",
"fred_md",
"hospital",
]
_MODEL_NAME = flags.DEFINE_string(
"model_name",
"timegpt-1-long-horizon",
"Path to model, can also be set to timegpt-1",
)
_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
QUANTILES = list(np.arange(1, 10) / 10.0)
def main():
results_list = []
run_id = np.random.randint(100000)
model_name = _MODEL_NAME.value
for dataset in dataset_names:
print(f"Evaluating model {model_name} on dataset {dataset}", flush=True)
exp = ExperimentHandler(dataset, quantiles=QUANTILES)
train_df = exp.train_df
horizon = exp.horizon
seasonality = exp.seasonality
freq = exp.freq
level = exp.level
fcsts_df, total_time, model_name = run_timegpt(
train_df=train_df,
horizon=exp.horizon,
model=model_name,
seasonality=seasonality,
freq=freq,
dataset=dataset,
level=level,
)
time_df = pd.DataFrame({"time": [total_time], "model": model_name})
fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name)
results = exp.evaluate_from_predictions(
models=[model_name], fcsts_df=fcsts_df, times_df=time_df
)
print(results, flush=True)
results_list.append(results)
results_full = pd.concat(results_list)
save_path = os.path.join(_SAVE_DIR.value, str(run_id))
print(f"Saving results to {save_path}", flush=True)
os.makedirs(save_path, exist_ok=True)
results_full.to_csv(f"{save_path}/results.csv")
if __name__ == "__main__":
FLAGS = flags.FLAGS
FLAGS(sys.argv)
main()
v1/experiments/extended_benchmarks/run_timesfm.py
+152
View File
@@ -0,0 +1,152 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Evaluation script for timesfm."""
import os
import sys
import time
from absl import flags
import numpy as np
import pandas as pd
import timesfm
from .utils import ExperimentHandler
dataset_names = [
"m1_monthly",
"m1_quarterly",
"m1_yearly",
"m3_monthly",
"m3_other",
"m3_quarterly",
"m3_yearly",
"m4_quarterly",
"m4_yearly",
"tourism_monthly",
"tourism_quarterly",
"tourism_yearly",
"nn5_daily_without_missing",
"m5",
"nn5_weekly",
"traffic",
"weather",
"australian_electricity_demand",
"car_parts_without_missing",
"cif_2016",
"covid_deaths",
"ercot",
"ett_small_15min",
"ett_small_1h",
"exchange_rate",
"fred_md",
"hospital",
]
context_dict_v2 = {}
context_dict_v1 = {
"cif_2016": 32,
"tourism_yearly": 64,
"covid_deaths": 64,
"tourism_quarterly": 64,
"tourism_monthly": 64,
"m1_monthly": 64,
"m1_quarterly": 64,
"m1_yearly": 64,
"m3_monthly": 64,
"m3_other": 64,
"m3_quarterly": 64,
"m3_yearly": 64,
"m4_quarterly": 64,
"m4_yearly": 64,
}
_MODEL_PATH = flags.DEFINE_string("model_path", "google/timesfm-2.0-500m-jax",
"Path to model")
_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, "Batch size")
_HORIZON = flags.DEFINE_integer("horizon", 128, "Horizon")
_BACKEND = flags.DEFINE_string("backend", "gpu", "Backend")
_NUM_JOBS = flags.DEFINE_integer("num_jobs", 1, "Number of jobs")
_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
QUANTILES = list(np.arange(1, 10) / 10.0)
def main():
results_list = []
model_path = _MODEL_PATH.value
num_layers = 20
max_context_len = 512
use_positional_embedding = True
context_dict = context_dict_v1
if "2.0" in model_path:
num_layers = 50
use_positional_embedding = False
max_context_len = 2048
context_dict = context_dict_v2
tfm = timesfm.TimesFm(
hparams=timesfm.TimesFmHparams(
backend="gpu",
per_core_batch_size=32,
horizon_len=128,
num_layers=num_layers,
context_len=max_context_len,
use_positional_embedding=use_positional_embedding,
),
checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=model_path),
)
run_id = np.random.randint(100000)
model_name = "timesfm"
for dataset in dataset_names:
print(f"Evaluating model {model_name} on dataset {dataset}", flush=True)
exp = ExperimentHandler(dataset, quantiles=QUANTILES)
if dataset in context_dict:
context_len = context_dict[dataset]
else:
context_len = max_context_len
train_df = exp.train_df
freq = exp.freq
init_time = time.time()
fcsts_df = tfm.forecast_on_df(
inputs=train_df,
freq=freq,
value_name="y",
model_name=model_name,
forecast_context_len=context_len,
num_jobs=_NUM_JOBS.value,
normalize=True,
)
total_time = time.time() - init_time
time_df = pd.DataFrame({"time": [total_time], "model": model_name})
results = exp.evaluate_from_predictions(models=[model_name],
fcsts_df=fcsts_df,
times_df=time_df)
print(results, flush=True)
results_list.append(results)
results_full = pd.concat(results_list)
save_path = os.path.join(_SAVE_DIR.value, str(run_id))
print(f"Saving results to {save_path}", flush=True)
os.makedirs(save_path, exist_ok=True)
results_full.to_csv(f"{save_path}/results.csv")
if __name__ == "__main__":
FLAGS = flags.FLAGS
FLAGS(sys.argv)
main()
v1/experiments/extended_benchmarks/tfm_extended_new.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 301 KiB

v1/experiments/extended_benchmarks/tfm_results.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 329 KiB

v1/experiments/extended_benchmarks/utils.py
+278
View File
@@ -0,0 +1,278 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py."""
from functools import partial
from itertools import repeat
import multiprocessing
import os
from pathlib import Path
from typing import List
from gluonts.dataset import Dataset
from gluonts.dataset.repository.datasets import (
dataset_names as gluonts_datasets,
get_dataset,
)
from gluonts.time_feature.seasonality import get_seasonality
import numpy as np
import pandas as pd
from utilsforecast.evaluation import evaluate
from utilsforecast.losses import mae, mase, smape
def parallel_transform(inp):
ts, last_n = inp[0], inp[1]
return ExperimentHandler._transform_gluonts_instance_to_df(ts, last_n=last_n)
def quantile_loss(
df: pd.DataFrame,
models: list,
q: float = 0.5,
id_col: str = "unique_id",
target_col: str = "y",
) -> pd.DataFrame:
delta_y = df[models].sub(df[target_col], axis=0)
res = (
np.maximum(q * delta_y, (q - 1) * delta_y)
.groupby(df[id_col], observed=True)
.mean()
)
res.index.name = id_col
res = res.reset_index()
return res
class ExperimentHandler:
def __init__(
self,
dataset: str,
quantiles: List[float] = list(np.arange(1, 10) / 10.0),
results_dir: str = "./results",
models_dir: str = "./models",
):
if dataset not in gluonts_datasets:
raise Exception(
f"dataset {dataset} not found in gluonts "
f"available datasets: {', '.join(gluonts_datasets)}"
)
self.dataset = dataset
self.quantiles = quantiles
self.level = self._transform_quantiles_to_levels(quantiles)
self.results_dir = results_dir
self.models_dir = models_dir
# defining datasets
self._maybe_download_m3_or_m5_file(self.dataset)
gluonts_dataset = get_dataset(self.dataset)
self.horizon = gluonts_dataset.metadata.prediction_length
if self.horizon is None:
raise Exception(
f"horizon not found for dataset {self.dataset} "
"experiment cannot be run"
)
self.freq = gluonts_dataset.metadata.freq
# get_seasonality() returns 1 for freq='D', override this to 7. This significantly improves the accuracy of
# statistical models on datasets like m5/nn5_daily. The models like AutoARIMA/AutoETS can still set
# seasonality=1 internally on datasets like weather by choosing non-seasonal models during model selection.
if self.freq == "D":
self.seasonality = 7
else:
self.seasonality = get_seasonality(self.freq)
self.gluonts_train_dataset = gluonts_dataset.train
self.gluonts_test_dataset = gluonts_dataset.test
self._create_dir_if_not_exists(self.results_dir)
try:
multiprocessing.set_start_method("spawn")
except RuntimeError:
print("Multiprocessing context has already been set.")
@staticmethod
def _maybe_download_m3_or_m5_file(dataset: str):
if dataset[:2] == "m3":
m3_file = Path.home() / ".gluonts" / "datasets" / "M3C.xls"
if not m3_file.exists():
from datasetsforecast.m3 import M3
from datasetsforecast.utils import download_file
download_file(m3_file.parent, M3.source_url)
elif dataset == "m5":
m5_raw_dir = Path.home() / ".gluonts" / "m5"
if not m5_raw_dir.exists():
import zipfile
from datasetsforecast.m5 import M5
from datasetsforecast.utils import download_file
download_file(m5_raw_dir, M5.source_url)
with zipfile.ZipFile(m5_raw_dir / "m5.zip", "r") as zip_ref:
zip_ref.extractall(m5_raw_dir)
@staticmethod
def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]:
level = [
int(100 - 200 * q) for q in quantiles if q < 0.5
] # in this case mean=mediain
level = sorted(list(set(level)))
return level
@staticmethod
def _create_dir_if_not_exists(directory: str):
Path(directory).mkdir(parents=True, exist_ok=True)
@staticmethod
def _transform_gluonts_instance_to_df(
ts: dict,
last_n: int | None = None,
) -> pd.DataFrame:
start_period = ts["start"]
start_ds, freq = start_period.to_timestamp(), start_period.freq
target = ts["target"]
ds = pd.date_range(start=start_ds, freq=freq, periods=len(target))
if last_n is not None:
target = target[-last_n:]
ds = ds[-last_n:]
ts_df = pd.DataFrame({"unique_id": ts["item_id"], "ds": ds, "y": target})
return ts_df
@staticmethod
def _transform_gluonts_dataset_to_df(
gluonts_dataset: Dataset,
last_n: int | None = None,
) -> pd.DataFrame:
with multiprocessing.Pool(os.cpu_count()) as pool: # Create a process pool
results = pool.map(
parallel_transform, zip(gluonts_dataset, repeat(last_n))
)
df = pd.concat(results)
df = df.reset_index(drop=True)
return df
@property
def train_df(self) -> pd.DataFrame:
train_df = self._transform_gluonts_dataset_to_df(self.gluonts_train_dataset)
return train_df
@property
def test_df(self) -> pd.DataFrame:
test_df = self._transform_gluonts_dataset_to_df(
self.gluonts_test_dataset,
last_n=self.horizon,
)
# Make sure that only the first backtest window is used for evaluation on `traffic` / `exchange_rate` datasets
return test_df.groupby("unique_id", sort=False).head(self.horizon)
def save_dataframe(self, df: pd.DataFrame, file_name: str):
df.to_csv(f"{self.results_dir}/{file_name}", index=False)
def save_results(
self, fcst_df: pd.DataFrame, total_time: float, model_name: str
):
self.save_dataframe(
fcst_df,
f"{model_name}-{self.dataset}-fcst.csv",
)
time_df = pd.DataFrame({"time": [total_time], "model": model_name})
self.save_dataframe(
time_df,
f"{model_name}-{self.dataset}-time.csv",
)
def fcst_from_level_to_quantiles(
self,
fcst_df: pd.DataFrame,
model_name: str,
) -> pd.DataFrame:
fcst_df = fcst_df.copy()
cols = ["unique_id", "ds", model_name]
for q in self.quantiles:
if q == 0.5:
col = f"{model_name}"
else:
lv = int(100 - 200 * q)
hi_or_lo = "lo" if lv > 0 else "hi"
lv = abs(lv)
col = f"{model_name}-{hi_or_lo}-{lv}"
q_col = f"{model_name}-q-{q}"
fcst_df[q_col] = fcst_df[col].values
cols.append(q_col)
return fcst_df[cols]
def evaluate_models(self, models: List[str]) -> pd.DataFrame:
fcsts_df = []
times_df = []
for model in models:
fcst_method_df = pd.read_csv(
f"{self.results_dir}/{model}-{self.dataset}-fcst.csv"
).set_index(["unique_id", "ds"])
fcsts_df.append(fcst_method_df)
time_method_df = pd.read_csv(
f"{self.results_dir}/{model}-{self.dataset}-time.csv"
)
times_df.append(time_method_df)
fcsts_df = pd.concat(fcsts_df, axis=1).reset_index()
fcsts_df["ds"] = pd.to_datetime(fcsts_df["ds"])
times_df = pd.concat(times_df)
return self.evaluate_from_predictions(
models=models, fcsts_df=fcsts_df, times_df=times_df
)
def evaluate_from_predictions(
self, models: List[str], fcsts_df: pd.DataFrame, times_df: pd.DataFrame
) -> pd.DataFrame:
test_df = self.test_df
train_df = self.train_df
test_df = test_df.merge(fcsts_df, how="left")
assert test_df.isna().sum().sum() == 0, "merge contains nas"
# point evaluation
point_fcsts_cols = ["unique_id", "ds", "y"] + models
test_df["unique_id"] = test_df["unique_id"].astype(str)
train_df["unique_id"] = train_df["unique_id"].astype(str)
mase_seas = partial(mase, seasonality=self.seasonality)
eval_df = evaluate(
test_df[point_fcsts_cols],
train_df=train_df,
metrics=[smape, mase_seas, mae],
)
# probabilistic evaluation
eval_prob_df = []
for q in self.quantiles:
prob_cols = [f"{model}-q-{q}" for model in models]
eval_q_df = quantile_loss(test_df, models=prob_cols, q=q)
eval_q_df[prob_cols] = eval_q_df[prob_cols] * self.horizon
eval_q_df = eval_q_df.rename(columns=dict(zip(prob_cols, models)))
eval_q_df["metric"] = f"quantile-loss-{q}"
eval_prob_df.append(eval_q_df)
eval_prob_df = pd.concat(eval_prob_df)
eval_prob_df = eval_prob_df.groupby("metric").sum().reset_index()
total_y = test_df["y"].sum()
eval_prob_df[models] = eval_prob_df[models] / total_y
eval_prob_df["metric"] = "scaled_crps"
eval_df = pd.concat([eval_df, eval_prob_df]).reset_index(drop=True)
eval_df = eval_df.groupby("metric").mean(numeric_only=True).reset_index()
eval_df = eval_df.melt(
id_vars="metric", value_name="value", var_name="model"
)
times_df.insert(0, "metric", "time")
times_df = times_df.rename(columns={"time": "value"})
eval_df = pd.concat([eval_df, times_df])
eval_df.insert(0, "dataset", self.dataset)
eval_df = eval_df.sort_values(["dataset", "metric", "model"])
eval_df = eval_df.reset_index(drop=True)
return eval_df
if __name__ == "__main__":
multiprocessing.set_start_method("spawn")
v1/experiments/long_horizon_benchmarks/README.md
+45
View File
@@ -0,0 +1,45 @@
# Extended Benchmarks
We benchmark on the original test set for ETT datasets as per long horizon benchmark papers (see [here](https://openreview.net/forum?id=pCbC3aQB5W) for example.) In the original benchmark, rolling validation task on all test windows (with a stride of 1) is considered. While we can easily run our method on this task, the baselines can take a very long time to run. Therefore we present results on a modified task with stride between windows set to Horizon length i.e all disjoint horizons in the test period is considered.
All experiments were performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). We compare TimesFM with [Amazon-Chronos](https://github.com/amazon-science/chronos-forecasting).
## Running TimesFM on the benchmark
We need to add the following packages for running these benchmarks. Follow the installation instructions till before `poetry lock`. Then,
```
poetry add git+https://github.com/awslabs/gluon-ts.git
poetry add git+https://github.com/amazon-science/chronos-forecasting.git
poetry lock
poetry install --only pax
```
Note that for now only the pax version runs on this benchmark, because we had to remove the old tf dependency from the pytorch version. We will fix this issue soon.
To run the timesfm on the benchmark do:
```
poetry run python3 -m experiments.long_horizon_benchmarks.run_eval \
--model_path=google/timesfm-1.0-200m --backend="gpu" \
--pred_len=96 --context_len=512 --dataset=etth1
```
In the above, `<model_path>` should point to the checkpoint directory that can be downloaded from HuggingFace.
For running chronos on the same benchmark you can run the command,
```
poetry run python3 -m experiments.long_horizon_benchmarks.run_eval \
--model_path=amazon/chronos-t5-mini --backend="gpu" \
--pred_len=96 --context_len=512 --dataset=etth1
```
You can change the model size from "mini" to "large" as required. The datasets we benchmark on are etth1, etth2, ettm1 and ettm2.
## Benchmark Results for TimesFM-1.0
![Benchmark Results Table](./tfm_long_horizon.png)
We compare the performance on horizon lengths of 96, 192 and 336, while context length is held fixed at 512.
We can see that TimesFM performs the best in terms of both wape and smape. More importantly it is much faster than the other methods, in particular it is more than 1000x faster than Chronos (Large).
v1/experiments/long_horizon_benchmarks/run_eval.py
+238
View File
@@ -0,0 +1,238 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Eval pipeline."""
import json
import os
import sys
import time
from absl import flags
import chronos
import numpy as np
import pandas as pd
import timesfm
from timesfm import data_loader
import torch
import tqdm
FLAGS = flags.FLAGS
_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64,
"Batch size for the randomly sampled batch")
_DATASET = flags.DEFINE_string("dataset", "etth1", "The name of the dataset.")
_MODEL_PATH = flags.DEFINE_string("model_path", "google/timesfm-2.0-500m-jax",
"The name of the model.")
_DATETIME_COL = flags.DEFINE_string("datetime_col", "date",
"Column having datetime.")
_NUM_COV_COLS = flags.DEFINE_list("num_cov_cols", None,
"Column having numerical features.")
_CAT_COV_COLS = flags.DEFINE_list("cat_cov_cols", None,
"Column having categorical features.")
_TS_COLS = flags.DEFINE_list("ts_cols", None, "Columns of time-series features")
_NORMALIZE = flags.DEFINE_bool("normalize", True,
"normalize data for eval or not")
_CONTEXT_LEN = flags.DEFINE_integer("context_len", 2048,
"Length of the context window")
_PRED_LEN = flags.DEFINE_integer("pred_len", 96, "prediction length.")
_BACKEND = flags.DEFINE_string("backend", "gpu", "backend to use")
_RESULTS_DIR = flags.DEFINE_string("results_dir", "./results/long_horizon",
"results directory")
DATA_DICT = {
"ettm2": {
"boundaries": [34560, 46080, 57600],
"data_path": "./datasets/ETT-small/ETTm2.csv",
"freq": "15min",
},
"ettm1": {
"boundaries": [34560, 46080, 57600],
"data_path": "./datasets/ETT-small/ETTm1.csv",
"freq": "15min",
},
"etth2": {
"boundaries": [8640, 11520, 14400],
"data_path": "./datasets/ETT-small/ETTh2.csv",
"freq": "H",
},
"etth1": {
"boundaries": [8640, 11520, 14400],
"data_path": "./datasets/ETT-small/ETTh1.csv",
"freq": "H",
},
"elec": {
"boundaries": [18413, 21044, 26304],
"data_path": "./datasets/electricity/electricity.csv",
"freq": "H",
},
"traffic": {
"boundaries": [12280, 14036, 17544],
"data_path": "./datasets/traffic/traffic.csv",
"freq": "H",
},
"weather": {
"boundaries": [36887, 42157, 52696],
"data_path": "./datasets/weather/weather.csv",
"freq": "10min",
},
}
QUANTILES = list(np.arange(1, 10) / 10.0)
EPS = 1e-7
def get_forecasts(model_path, model, past, freq, pred_len):
"""Get forecasts."""
if model_path.startswith("amazon"):
out = model.predict(
torch.tensor(past),
prediction_length=pred_len,
limit_prediction_length=False,
)
out = out.numpy()
out = np.median(out, axis=1)
else:
lfreq = [freq] * past.shape[0]
_, out = model.forecast(list(past), lfreq)
out = out[:, :, 5]
return out
def _mse(y_pred, y_true):
"""mse loss."""
return np.square(y_pred - y_true)
def _mae(y_pred, y_true):
"""mae loss."""
return np.abs(y_pred - y_true)
def _smape(y_pred, y_true):
"""_smape loss."""
abs_diff = np.abs(y_pred - y_true)
abs_val = (np.abs(y_true) + np.abs(y_pred)) / 2
abs_val = np.where(abs_val > EPS, abs_val, 1.0)
abs_diff = np.where(abs_val > EPS, abs_diff, 0.0)
return abs_diff / abs_val
def eval():
"""Eval pipeline."""
dataset = _DATASET.value
data_path = DATA_DICT[dataset]["data_path"]
freq = DATA_DICT[dataset]["freq"]
int_freq = timesfm.freq_map(freq)
boundaries = DATA_DICT[dataset]["boundaries"]
data_df = pd.read_csv(open(data_path, "r"))
if _TS_COLS.value is not None:
ts_cols = DATA_DICT[dataset]["ts_cols"]
num_cov_cols = DATA_DICT[dataset]["num_cov_cols"]
cat_cov_cols = DATA_DICT[dataset]["cat_cov_cols"]
else:
ts_cols = [col for col in data_df.columns if col != _DATETIME_COL.value]
num_cov_cols = None
cat_cov_cols = None
batch_size = min(_BATCH_SIZE.value, len(ts_cols))
dtl = data_loader.TimeSeriesdata(
data_path=data_path,
datetime_col=_DATETIME_COL.value,
num_cov_cols=num_cov_cols,
cat_cov_cols=cat_cov_cols,
ts_cols=np.array(ts_cols),
train_range=[0, boundaries[0]],
val_range=[boundaries[0], boundaries[1]],
test_range=[boundaries[1], boundaries[2]],
hist_len=_CONTEXT_LEN.value,
pred_len=_PRED_LEN.value,
batch_size=batch_size,
freq=freq,
normalize=_NORMALIZE.value,
epoch_len=None,
holiday=False,
permute=False,
)
eval_itr = dtl.tf_dataset(mode="test",
shift=_PRED_LEN.value).as_numpy_iterator()
model_path = _MODEL_PATH.value
if model_path.startswith("amazon"):
model = chronos.ChronosPipeline.from_pretrained(
model_path,
device_map="auto",
torch_dtype=torch.bfloat16,
)
else:
model = timesfm.TimesFm(
hparams=timesfm.TimesFmHparams(
backend="gpu",
per_core_batch_size=32,
horizon_len=128,
num_layers=50,
context_len=_CONTEXT_LEN.value,
use_positional_embedding=False,
),
checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=model_path),
)
smape_run_losses = []
mse_run_losses = []
mae_run_losses = []
num_elements = 0
abs_sum = 0
start_time = time.time()
for batch in tqdm.tqdm(eval_itr):
past = batch[0]
actuals = batch[3]
forecasts = get_forecasts(model_path, model, past, int_freq,
_PRED_LEN.value)
forecasts = forecasts[:, 0:actuals.shape[1]]
mae_run_losses.append(_mae(forecasts, actuals).sum())
mse_run_losses.append(_mse(forecasts, actuals).sum())
smape_run_losses.append(_smape(forecasts, actuals).sum())
num_elements += actuals.shape[0] * actuals.shape[1]
abs_sum += np.abs(actuals).sum()
mse_val = np.sum(mse_run_losses) / num_elements
result_dict = {
"mse": mse_val,
"smape": np.sum(smape_run_losses) / num_elements,
"mae": np.sum(mae_run_losses) / num_elements,
"wape": np.sum(mae_run_losses) / abs_sum,
"nrmse": np.sqrt(mse_val) / (abs_sum / num_elements),
"num_elements": num_elements,
"abs_sum": abs_sum,
"total_time": time.time() - start_time,
"model_path": model_path,
"dataset": dataset,
"freq": freq,
"pred_len": _PRED_LEN.value,
"context_len": _CONTEXT_LEN.value,
}
run_id = np.random.randint(10000)
save_path = os.path.join(_RESULTS_DIR.value, str(run_id))
print(f"Saving results to {save_path}", flush=True)
os.makedirs(save_path, exist_ok=True)
with open(os.path.join(save_path, "results.json"), "w") as f:
json.dump(result_dict, f)
print(result_dict, flush=True)
if __name__ == "__main__":
FLAGS = flags.FLAGS
FLAGS(sys.argv)
eval()
v1/experiments/long_horizon_benchmarks/tfm_long_horizon.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 193 KiB

v1/notebooks/covariates.ipynb
+383
View File
@@ -0,0 +1,383 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# TimesFM with Covariates\n",
"\n",
"This toturial notebook demonstrates how to utilize exogenous covariates with TimesFM when making forecasts. Before running this notebook, make sure:\n",
"\n",
"- You've read through the README of TimesFM.\n",
"- A local kernel with Python 3.10 is up and running, for the jax version.\n",
"- Install the JAX version following the installation instructions."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup the environment and install TimesFM."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load the checkpoint\n",
"\n",
"**Notice:** Please set up the backend as per your machine (\"cpu\", \"gpu\" or \"tpu\"). This notebook will run by default on GPU.\n",
"\n",
"We load the 2.0-500m model checkpoint from HuggingFace."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import timesfm\n",
"timesfm_backend = \"gpu\" # @param\n",
"\n",
"model = timesfm.TimesFm(\n",
" hparams=timesfm.TimesFmHparams(\n",
" backend=timesfm_backend,\n",
" per_core_batch_size=32,\n",
" horizon_len=128,\n",
" num_layers=50,\n",
" use_positional_embedding=False,\n",
" context_len=2048,\n",
" ),\n",
" checkpoint=timesfm.TimesFmCheckpoint(\n",
" huggingface_repo_id=\"google/timesfm-2.0-500m-jax\"),\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Covariates\n",
"\n",
"Let's take a toy example of forecasting sales for a grocery store: \n",
"\n",
"**Task:** Given the observed the daily sales of this week (7 days), forecast the daily sales of next week (7 days).\n",
"\n",
"```\n",
"Product: ice cream\n",
"Daily_sales: [30, 30, 4, 5, 7, 8, 10]\n",
"Category: food\n",
"Base_price: 1.99\n",
"Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]\n",
"Has_promotion: [Yes, Yes, No, No, No, Yes, Yes, No, No, No, No, No, No, No]\n",
"Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]\n",
"```\n",
"\n",
"```\n",
"Product: sunscreen\n",
"Daily_sales: [5, 7, 12, 13, 5, 6, 10]\n",
"Category: skin product\n",
"Base_price: 29.99\n",
"Weekday: [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]\n",
"Has_promotion: [No, No, Yes, Yes, No, No, No, Yes, Yes, Yes, Yes, Yes, Yes, Yes]\n",
"Daily_temperature: [31.0, 24.3, 19.4, 26.2, 24.6, 30.0, 31.1, 32.4, 30.9, 26.0, 25.0, 27.8, 29.5, 31.2]\n",
"```\n",
"\n",
"In this example, besides the `Daily_sales`, we also have covariates `Category`, `Base_price`, `Weekday`, `Has_promotion`, `Daily_temperature`. Let's introduce some concepts:\n",
"\n",
"**Static covariates** are covariates for each time series. \n",
"- In our example, `Category` is a **static categorical covariate**, \n",
"- `Base_price` is a **static numerical covariates**.\n",
"\n",
"**Dynamic covariates** are covaraites for each time stamps.\n",
"- Date / time related features can be usually treated as dynamic covariates.\n",
"- In our example, `Weekday` and `Has_promotion` are **dynamic categorical covariates**.\n",
"- `Daily_temperate` is a **dynamic numerical covariate**.\n",
"\n",
"**Notice:** Here we make it mandatory that the dynamic covariates need to cover both the forecasting context and horizon. For example, all dynamic covariates in the example have 14 values: the first 7 correspond to the observed 7 days, and the last 7 correspond to the next 7 days."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# TimesFM with Covariates\n",
"\n",
"\n",
"The strategy we take here is to treat covariates as batched in-context exogenous regressors (XReg) and fit linear models on them outside of TimesFM. The final forecast will be the sum of the TimesFM forecast and the linear model forecast.\n",
"\n",
" In simple words, we consider these two options.\n",
"\n",
"**Option 1:** Get the TimesFM forecast, and fit the linear model regressing the residuals on the covariates (\"timesfm + xreg\").\n",
"\n",
"**Option 2:** Fit the linear model of the time series itself on the covariates, then forecast the residuals using TimesFM (\"xreg + timesfm\").\n",
"\n",
"Let's take a code at the example of Electricity Price Forecasting (EPF). \n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import pandas as pd\n",
"import numpy as np\n",
"from collections import defaultdict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"df = pd.read_csv('https://datasets-nixtla.s3.amazonaws.com/EPF_FR_BE.csv')\n",
"df['ds'] = pd.to_datetime(df['ds'])\n",
"df"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This dataset has a few covariates beside the hourly target `y`:\n",
"\n",
"- `unique_id`: a static categorical covariate indicating the country.\n",
"- `gen_forecast`: a dynamic numerical covariate indicating the estimated electricity to be generated.\n",
"- `system_load`: the observed system load. Notice that this **CANNOT** be considered as a dynamic numerical covariate because we cannot know its values over the forecasting horizon in advance.\n",
"- `weekday`: a dynamic categorical covariate.\\\n",
"\n",
"Let's now make some forecasting tasks for TimesFM based on this dataset. For simplicity we create forecast contexts of 120 time points (hours) and forecast horizons of 24 time points."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"# Data pipelining\n",
"def get_batched_data_fn(\n",
" batch_size: int = 128, \n",
" context_len: int = 120, \n",
" horizon_len: int = 24,\n",
"):\n",
" examples = defaultdict(list)\n",
"\n",
" num_examples = 0\n",
" for country in (\"FR\", \"BE\"):\n",
" sub_df = df[df[\"unique_id\"] == country]\n",
" for start in range(0, len(sub_df) - (context_len + horizon_len), horizon_len):\n",
" num_examples += 1\n",
" examples[\"country\"].append(country)\n",
" examples[\"inputs\"].append(sub_df[\"y\"][start:(context_end := start + context_len)].tolist())\n",
" examples[\"gen_forecast\"].append(sub_df[\"gen_forecast\"][start:context_end + horizon_len].tolist())\n",
" examples[\"week_day\"].append(sub_df[\"week_day\"][start:context_end + horizon_len].tolist())\n",
" examples[\"outputs\"].append(sub_df[\"y\"][context_end:(context_end + horizon_len)].tolist())\n",
" \n",
" def data_fn():\n",
" for i in range(1 + (num_examples - 1) // batch_size):\n",
" yield {k: v[(i * batch_size) : ((i + 1) * batch_size)] for k, v in examples.items()}\n",
" \n",
" return data_fn\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"# Define metrics\n",
"def mse(y_pred, y_true):\n",
" y_pred = np.array(y_pred)\n",
" y_true = np.array(y_true)\n",
" return np.mean(np.square(y_pred - y_true), axis=1, keepdims=True)\n",
"\n",
"def mae(y_pred, y_true):\n",
" y_pred = np.array(y_pred)\n",
" y_true = np.array(y_true)\n",
" return np.mean(np.abs(y_pred - y_true), axis=1, keepdims=True)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now let's try `model.forecast_with_covariates`. \n",
"\n",
"In particular, the output is a tuple whose first element is the new forecast."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import time\n",
"\n",
"# Benchmark\n",
"batch_size = 128\n",
"context_len = 120\n",
"horizon_len = 24\n",
"input_data = get_batched_data_fn(batch_size = 128)\n",
"metrics = defaultdict(list)\n",
"\n",
"\n",
"for i, example in enumerate(input_data()):\n",
" raw_forecast, _ = model.forecast(\n",
" inputs=example[\"inputs\"], freq=[0] * len(example[\"inputs\"])\n",
" )\n",
" start_time = time.time()\n",
" # Forecast with covariates\n",
" # Output: new forecast, forecast by the xreg\n",
" cov_forecast, ols_forecast = model.forecast_with_covariates( \n",
" inputs=example[\"inputs\"],\n",
" dynamic_numerical_covariates={\n",
" \"gen_forecast\": example[\"gen_forecast\"],\n",
" },\n",
" dynamic_categorical_covariates={\n",
" \"week_day\": example[\"week_day\"],\n",
" },\n",
" static_numerical_covariates={},\n",
" static_categorical_covariates={\n",
" \"country\": example[\"country\"]\n",
" },\n",
" freq=[0] * len(example[\"inputs\"]),\n",
" xreg_mode=\"xreg + timesfm\", # default\n",
" ridge=0.0,\n",
" force_on_cpu=False,\n",
" normalize_xreg_target_per_input=True, # default\n",
" )\n",
" print(\n",
" f\"\\rFinished batch {i} linear in {time.time() - start_time} seconds\",\n",
" end=\"\",\n",
" )\n",
" metrics[\"eval_mae_timesfm\"].extend(\n",
" mae(raw_forecast[:, :horizon_len], example[\"outputs\"])\n",
" )\n",
" metrics[\"eval_mae_xreg_timesfm\"].extend(mae(cov_forecast, example[\"outputs\"]))\n",
" metrics[\"eval_mae_xreg\"].extend(mae(ols_forecast, example[\"outputs\"]))\n",
" metrics[\"eval_mse_timesfm\"].extend(\n",
" mse(raw_forecast[:, :horizon_len], example[\"outputs\"])\n",
" )\n",
" metrics[\"eval_mse_xreg_timesfm\"].extend(mse(cov_forecast, example[\"outputs\"]))\n",
" metrics[\"eval_mse_xreg\"].extend(mse(ols_forecast, example[\"outputs\"]))\n",
"\n",
"print()\n",
"\n",
"for k, v in metrics.items():\n",
" print(f\"{k}: {np.mean(v)}\")\n",
"\n",
"# My output:\n",
"# eval_mae_timesfm: 6.762283045916956\n",
"# eval_mae_xreg_timesfm: 5.39219617611074\n",
"# eval_mae_xreg: 37.15275842572484\n",
"# eval_mse_timesfm: 166.7771466306823\n",
"# eval_mse_xreg_timesfm: 120.64757721021306\n",
"# eval_mse_xreg: 1672.2116821201796"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You should see results close to \n",
"```\n",
"eval_mae_timesfm: 6.729583250571446\n",
"eval_mae_xreg_timesfm: 5.3375301110158\n",
"eval_mae_xreg: 37.152760709266\n",
"eval_mse_timesfm: 162.3132151851567\n",
"eval_mse_xreg_timesfm: 120.9900627409689\n",
"eval_mse_xreg: 1672.208769045399\n",
"```\n",
"\n",
"With the covariates, the TimesFM forecast Mean Absolute Error improves from 6.73 to 5.34, and Mean Squred Error from 162.31 to 120.99. The results of purely fitting the linear model are also provided for reference."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Formatting Your Request\n",
"\n",
"It is quite crucial to get the covariates properly formatted so that we can call this `model.forecast_with_covariates`. Please see its docstring for details. Here let's also grab a batch from a toy data input pipeline for quick explanations."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"toy_input_pipeline = get_batched_data_fn(batch_size=2, context_len=5, horizon_len=2)\n",
"print(next(toy_input_pipeline()))\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You should see something similar to this\n",
"```\n",
"{\n",
" 'country': ['FR', 'FR'], \n",
" 'inputs': [[53.48, 51.93, 48.76, 42.27, 38.41], [48.76, 42.27, 38.41, 35.72, 32.66]], \n",
" 'gen_forecast': [[76905.0, 75492.0, 74394.0, 72639.0, 69347.0, 67960.0, 67564.0], [74394.0, 72639.0, 69347.0, 67960.0, 67564.0, 67277.0, 67019.0]], \n",
" 'week_day': [[3, 3, 3, 3, 3, 3, 3], [3, 3, 3, 3, 3, 3, 3]], \n",
" 'outputs': [[35.72, 32.66], [32.83, 30.06]],\n",
"}\n",
"```\n",
"\n",
"Notice:\n",
"- We have two examples in this batch.\n",
"- For each example we support different context lengths and horizon lengths just as `model.forecast`. Although it is not demonstrated in this dataset.\n",
"- If dynamic covariates are present, the horizon lengths will be inferred from them, e.g. how many values are provided in additional to the ones corresponding to the inputs. Make sure all your dynamic covariates have the same length per example.\n",
"- The static covariates are one per example.\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## More Applications\n",
"\n",
"### Past Dynamic Covariates\n",
"\n",
"Past dynamic covariates are covariates that are only available for the context. For instance in our example `system_load` is a past dynamic covariate. Time series models generally can handle this, however it is something the batched in context regression cannot address, because these regressors are not available in the future. If you do have those covariates and consider them very meaningful, there are two hacky options to try immediately:\n",
"\n",
"1. Shift and repeat these past dynamic covariates to use their delayed version. For example, if you think the `system_load` for this week is meaningful for forecasting next week, you can create a `delay_7_system_load` by shifting 7 timestamps and use this as one dynamic numerical covariate for TimesFM.\n",
"2. Bootstrap, that is to run TimesFM once to forecast these past dynamic covariates into the horizon, then call TimesFM again using these forecasts as the future part for these dynamic covariates.\n",
"\n",
"### Multivariate Time Series\n",
"\n",
"For multivariate time series, if we need univariate forecast, we can try treating the main time series as the target and use the rest as the dynamic covariates."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "chronos-v2",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
v1/notebooks/finetuning.ipynb
+619
View File
@@ -0,0 +1,619 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Importing relevant packages for finetuning"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'\n",
"os.environ['JAX_PMAP_USE_TENSORSTORE'] = 'false'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import timesfm\n",
"import gc\n",
"import numpy as np\n",
"import pandas as pd\n",
"from timesfm import patched_decoder\n",
"from timesfm import data_loader"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"from tqdm import tqdm\n",
"import dataclasses\n",
"import IPython\n",
"import IPython.display\n",
"import matplotlib as mpl\n",
"import matplotlib.pyplot as plt\n",
"mpl.rcParams['figure.figsize'] = (8, 6)\n",
"mpl.rcParams['axes.grid'] = False"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Loading TimesFM pretrained checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"timesfm_backend = \"gpu\" # @param\n",
"\n",
"tfm = timesfm.TimesFm(\n",
" hparams=timesfm.TimesFmHparams(\n",
" backend=timesfm_backend,\n",
" per_core_batch_size=32,\n",
" horizon_len=128,\n",
" num_layers=50,\n",
" # Se this to True for v1.0 checkpoints\n",
" use_positional_embedding=False,\n",
" # Note that we could set this to as high as 2048 but keeping it 512 here so that\n",
" # both v1.0 and 2.0 checkpoints work\n",
" context_len=512,\n",
" ),\n",
" checkpoint=timesfm.TimesFmCheckpoint(\n",
" huggingface_repo_id=\"google/timesfm-2.0-500m-jax\"),\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluating pretrained checkpoint on ETT datasets"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"DATA_DICT = {\n",
" \"ettm2\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"ettm1\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"etth2\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"etth1\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"elec\": {\n",
" \"boundaries\": [18413, 21044, 26304],\n",
" \"data_path\": \"../datasets/electricity/electricity.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"traffic\": {\n",
" \"boundaries\": [12280, 14036, 17544],\n",
" \"data_path\": \"../datasets/traffic/traffic.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"weather\": {\n",
" \"boundaries\": [36887, 42157, 52696],\n",
" \"data_path\": \"../datasets/weather/weather.csv\",\n",
" \"freq\": \"10min\",\n",
" },\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"dataset = \"ettm1\"\n",
"data_path = DATA_DICT[dataset][\"data_path\"]\n",
"freq = DATA_DICT[dataset][\"freq\"]\n",
"int_freq = timesfm.freq_map(freq)\n",
"boundaries = DATA_DICT[dataset][\"boundaries\"]\n",
"\n",
"data_df = pd.read_csv(open(data_path, \"r\"))\n",
"\n",
"\n",
"ts_cols = [col for col in data_df.columns if col != \"date\"]\n",
"num_cov_cols = None\n",
"cat_cov_cols = None\n",
"\n",
"context_len = 512\n",
"pred_len = 96\n",
"\n",
"num_ts = len(ts_cols)\n",
"batch_size = 8\n",
"\n",
"dtl = data_loader.TimeSeriesdata(\n",
" data_path=data_path,\n",
" datetime_col=\"date\",\n",
" num_cov_cols=num_cov_cols,\n",
" cat_cov_cols=cat_cov_cols,\n",
" ts_cols=np.array(ts_cols),\n",
" train_range=[0, boundaries[0]],\n",
" val_range=[boundaries[0], boundaries[1]],\n",
" test_range=[boundaries[1], boundaries[2]],\n",
" hist_len=context_len,\n",
" pred_len=pred_len,\n",
" batch_size=num_ts,\n",
" freq=freq,\n",
" normalize=True,\n",
" epoch_len=None,\n",
" holiday=False,\n",
" permute=True,\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_batches = dtl.tf_dataset(mode=\"train\", shift=1).batch(batch_size)\n",
"val_batches = dtl.tf_dataset(mode=\"val\", shift=pred_len)\n",
"test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for tbatch in tqdm(train_batches.as_numpy_iterator()):\n",
" break\n",
"print(tbatch[0].shape)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### MAE on the test split for the pretrained TimesFM model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" forecasts, _ = tfm.forecast(list(past), [0] * past.shape[0], normalize=True)\n",
" forecasts = forecasts[:, 0 : actuals.shape[1]]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Finetuning the model on the ETT dataset"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"import jax\n",
"from jax import numpy as jnp\n",
"from praxis import pax_fiddle\n",
"from praxis import py_utils\n",
"from praxis import pytypes\n",
"from praxis import base_model\n",
"from praxis import optimizers\n",
"from praxis import schedules\n",
"from praxis import base_hyperparams\n",
"from praxis import base_layer\n",
"from paxml import tasks_lib\n",
"from paxml import trainer_lib\n",
"from paxml import checkpoints\n",
"from paxml import learners\n",
"from paxml import partitioning\n",
"from paxml import checkpoint_types"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"# PAX shortcuts\n",
"NestedMap = py_utils.NestedMap\n",
"WeightInit = base_layer.WeightInit\n",
"WeightHParams = base_layer.WeightHParams\n",
"InstantiableParams = py_utils.InstantiableParams\n",
"JTensor = pytypes.JTensor\n",
"NpTensor = pytypes.NpTensor\n",
"WeightedScalars = pytypes.WeightedScalars\n",
"instantiate = base_hyperparams.instantiate\n",
"LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]\n",
"AuxLossStruct = base_layer.AuxLossStruct\n",
"\n",
"AUX_LOSS = base_layer.AUX_LOSS\n",
"template_field = base_layer.template_field\n",
"\n",
"# Standard prng key names\n",
"PARAMS = base_layer.PARAMS\n",
"RANDOM = base_layer.RANDOM\n",
"\n",
"key = jax.random.PRNGKey(seed=1234)"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"model = pax_fiddle.Config(\n",
" patched_decoder.PatchedDecoderFinetuneModel,\n",
" name='patched_decoder_finetune',\n",
" core_layer_tpl=tfm.model_p,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### We will hold the transformer layers fixed while finetuning, while training all other components."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"@pax_fiddle.auto_config\n",
"def build_learner() -> learners.Learner:\n",
" return pax_fiddle.Config(\n",
" learners.Learner,\n",
" name='learner',\n",
" loss_name='avg_qloss',\n",
" optimizer=optimizers.Adam(\n",
" epsilon=1e-7,\n",
" clip_threshold=1e2,\n",
" learning_rate=1e-2,\n",
" lr_schedule=pax_fiddle.Config(\n",
" schedules.Cosine,\n",
" initial_value=1e-3,\n",
" final_value=1e-4,\n",
" total_steps=40000,\n",
" ),\n",
" ema_decay=0.9999,\n",
" ),\n",
" # Linear probing i.e we hold the transformer layers fixed.\n",
" bprop_variable_exclusion=['.*/stacked_transformer_layer/.*'],\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"task_p = tasks_lib.SingleTask(\n",
" name='ts-learn',\n",
" model=model,\n",
" train=tasks_lib.SingleTask.Train(\n",
" learner=build_learner(),\n",
" ),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"task_p.model.ici_mesh_shape = [1, 1, 1]\n",
"task_p.model.mesh_axis_names = ['replica', 'data', 'mdl']\n",
"\n",
"DEVICES = np.array(jax.devices()).reshape([1, 1, 1])\n",
"MESH = jax.sharding.Mesh(DEVICES, ['replica', 'data', 'mdl'])\n",
"\n",
"num_devices = jax.local_device_count()\n",
"print(f'num_devices: {num_devices}')\n",
"print(f'device kind: {jax.local_devices()[0].device_kind}')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"jax_task = task_p\n",
"key, init_key = jax.random.split(key)\n",
"\n",
"# To correctly prepare a batch of data for model initialization (now that shape\n",
"# inference is merged), we take one devices*batch_size tensor tuple of data,\n",
"# slice out just one batch, then run the prepare_input_batch function over it.\n",
"\n",
"\n",
"def process_train_batch(batch):\n",
" past_ts = batch[0].reshape(batch_size * num_ts, -1)\n",
" actual_ts = batch[3].reshape(batch_size * num_ts, -1)\n",
" return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n",
"\n",
"\n",
"def process_eval_batch(batch):\n",
" past_ts = batch[0]\n",
" actual_ts = batch[3]\n",
" return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n",
"\n",
"\n",
"jax_model_states, _ = trainer_lib.initialize_model_state(\n",
" jax_task,\n",
" init_key,\n",
" process_train_batch(tbatch),\n",
" checkpoint_type=checkpoint_types.CheckpointType.GDA,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Setting the initial model weights to the pretrained TimesFM parameters."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"jax_model_states.mdl_vars['params']['core_layer'] = tfm._train_state.mdl_vars['params']\n",
"jax_vars = jax_model_states.mdl_vars\n",
"gc.collect()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Training loop"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [],
"source": [
"jax_task = task_p\n",
"\n",
"\n",
"def train_step(states, prng_key, inputs):\n",
" return trainer_lib.train_step_single_learner(\n",
" jax_task, states, prng_key, inputs\n",
" )\n",
"\n",
"\n",
"def eval_step(states, prng_key, inputs):\n",
" states = states.to_eval_state()\n",
" return trainer_lib.eval_step_single_learner(\n",
" jax_task, states, prng_key, inputs\n",
" )\n",
"\n",
"key, train_key, eval_key = jax.random.split(key, 3)\n",
"train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())\n",
"eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())\n",
"\n",
"p_train_step = jax.pmap(train_step, axis_name='batch')\n",
"p_eval_step = jax.pmap(eval_step, axis_name='batch')"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [],
"source": [
"replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)\n",
"replicated_jax_vars = replicated_jax_states.mdl_vars"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [],
"source": [
"best_eval_loss = 1e7\n",
"step_count = 0\n",
"patience = 0\n",
"NUM_EPOCHS = 100\n",
"PATIENCE = 5\n",
"TRAIN_STEPS_PER_EVAL = 1000\n",
"CHECKPOINT_DIR='/home/senrajat_google_com/ettm1_finetune'"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
"def reshape_batch_for_pmap(batch, num_devices):\n",
" def _reshape(input_tensor):\n",
" bsize = input_tensor.shape[0]\n",
" residual_shape = list(input_tensor.shape[1:])\n",
" nbsize = bsize // num_devices\n",
" return jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)\n",
"\n",
" return jax.tree.map(_reshape, batch)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for epoch in range(NUM_EPOCHS):\n",
" print(f\"__________________Epoch: {epoch}__________________\", flush=True)\n",
" train_its = train_batches.as_numpy_iterator()\n",
" if patience >= PATIENCE:\n",
" print(\"Early stopping.\", flush=True)\n",
" break\n",
" for batch in tqdm(train_its):\n",
" train_losses = []\n",
" if patience >= PATIENCE:\n",
" print(\"Early stopping.\", flush=True)\n",
" break\n",
" tbatch = process_train_batch(batch)\n",
" tbatch = reshape_batch_for_pmap(tbatch, num_devices)\n",
" replicated_jax_states, step_fun_out = p_train_step(\n",
" replicated_jax_states, train_prng_seed, tbatch\n",
" )\n",
" train_losses.append(step_fun_out.loss[0])\n",
" if step_count % TRAIN_STEPS_PER_EVAL == 0:\n",
" print(\n",
" f\"Train loss at step {step_count}: {np.mean(train_losses)}\",\n",
" flush=True,\n",
" )\n",
" train_losses = []\n",
" print(\"Starting eval.\", flush=True)\n",
" val_its = val_batches.as_numpy_iterator()\n",
" eval_losses = []\n",
" for ev_batch in tqdm(val_its):\n",
" ebatch = process_eval_batch(ev_batch)\n",
" ebatch = reshape_batch_for_pmap(ebatch, num_devices)\n",
" _, step_fun_out = p_eval_step(\n",
" replicated_jax_states, eval_prng_seed, ebatch\n",
" )\n",
" eval_losses.append(step_fun_out.loss[0])\n",
" mean_loss = np.mean(eval_losses)\n",
" print(f\"Eval loss at step {step_count}: {mean_loss}\", flush=True)\n",
" if mean_loss < best_eval_loss or np.isnan(mean_loss):\n",
" best_eval_loss = mean_loss\n",
" print(\"Saving checkpoint.\")\n",
" jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(\n",
" replicated_jax_states\n",
" )\n",
" checkpoints.save_checkpoint(\n",
" jax_state_for_saving, CHECKPOINT_DIR, overwrite=True\n",
" )\n",
" patience = 0\n",
" del jax_state_for_saving\n",
" gc.collect()\n",
" else:\n",
" patience += 1\n",
" print(f\"patience: {patience}\")\n",
" step_count += 1"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Loading and evaluating the best (according to validation loss) finetuned checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_state = checkpoints.restore_checkpoint(jax_model_states, CHECKPOINT_DIR)\n",
"print(train_state.step)\n",
"tfm._train_state.mdl_vars['params'] = train_state.mdl_vars['params']['core_layer']\n",
"tfm.jit_decode()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
" forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## There is around a __7%__ reduction in MAE from finetuning."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "chronos-v2",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
v1/notebooks/finetuning_torch.ipynb
+538
View File
@@ -0,0 +1,538 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Introduction\n",
"This notebook shows how to use TimesFM with finetuning. \n",
"\n",
"In order to perform finetuning, you need to create the Pytorch Dataset in a proper format. The example of the Dataset is provided below.\n",
"The finetuning code can be found in timesfm.finetuning_torch.py. This notebook just imports the methods from finetuning"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Dataset Creation"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"TimesFM v1.2.0. See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs.\n",
"Loaded Jax TimesFM.\n",
"Loaded PyTorch TimesFM.\n"
]
}
],
"source": [
"from os import path\n",
"from typing import Optional, Tuple\n",
"\n",
"import numpy as np\n",
"import pandas as pd\n",
"import torch\n",
"import torch.multiprocessing as mp\n",
"import yfinance as yf\n",
"from finetuning.finetuning_torch import FinetuningConfig, TimesFMFinetuner\n",
"from huggingface_hub import snapshot_download\n",
"from torch.utils.data import Dataset\n",
"\n",
"from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams\n",
"from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder\n",
"import os\n",
"\n",
"\n",
"class TimeSeriesDataset(Dataset):\n",
" \"\"\"Dataset for time series data compatible with TimesFM.\"\"\"\n",
"\n",
" def __init__(self,\n",
" series: np.ndarray,\n",
" context_length: int,\n",
" horizon_length: int,\n",
" freq_type: int = 0):\n",
" \"\"\"\n",
" Initialize dataset.\n",
"\n",
" Args:\n",
" series: Time series data\n",
" context_length: Number of past timesteps to use as input\n",
" horizon_length: Number of future timesteps to predict\n",
" freq_type: Frequency type (0, 1, or 2)\n",
" \"\"\"\n",
" if freq_type not in [0, 1, 2]:\n",
" raise ValueError(\"freq_type must be 0, 1, or 2\")\n",
"\n",
" self.series = series\n",
" self.context_length = context_length\n",
" self.horizon_length = horizon_length\n",
" self.freq_type = freq_type\n",
" self._prepare_samples()\n",
"\n",
" def _prepare_samples(self) -> None:\n",
" \"\"\"Prepare sliding window samples from the time series.\"\"\"\n",
" self.samples = []\n",
" total_length = self.context_length + self.horizon_length\n",
"\n",
" for start_idx in range(0, len(self.series) - total_length + 1):\n",
" end_idx = start_idx + self.context_length\n",
" x_context = self.series[start_idx:end_idx]\n",
" x_future = self.series[end_idx:end_idx + self.horizon_length]\n",
" self.samples.append((x_context, x_future))\n",
"\n",
" def __len__(self) -> int:\n",
" return len(self.samples)\n",
"\n",
" def __getitem__(\n",
" self, index: int\n",
" ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
" x_context, x_future = self.samples[index]\n",
"\n",
" x_context = torch.tensor(x_context, dtype=torch.float32)\n",
" x_future = torch.tensor(x_future, dtype=torch.float32)\n",
"\n",
" input_padding = torch.zeros_like(x_context)\n",
" freq = torch.tensor([self.freq_type], dtype=torch.long)\n",
"\n",
" return x_context, input_padding, freq, x_future\n",
"\n",
"def prepare_datasets(series: np.ndarray,\n",
" context_length: int,\n",
" horizon_length: int,\n",
" freq_type: int = 0,\n",
" train_split: float = 0.8) -> Tuple[Dataset, Dataset]:\n",
" \"\"\"\n",
" Prepare training and validation datasets from time series data.\n",
"\n",
" Args:\n",
" series: Input time series data\n",
" context_length: Number of past timesteps to use\n",
" horizon_length: Number of future timesteps to predict\n",
" freq_type: Frequency type (0, 1, or 2)\n",
" train_split: Fraction of data to use for training\n",
"\n",
" Returns:\n",
" Tuple of (train_dataset, val_dataset)\n",
" \"\"\"\n",
" train_size = int(len(series) * train_split)\n",
" train_data = series[:train_size]\n",
" val_data = series[train_size:]\n",
"\n",
" # Create datasets with specified frequency type\n",
" train_dataset = TimeSeriesDataset(train_data,\n",
" context_length=context_length,\n",
" horizon_length=horizon_length,\n",
" freq_type=freq_type)\n",
"\n",
" val_dataset = TimeSeriesDataset(val_data,\n",
" context_length=context_length,\n",
" horizon_length=horizon_length,\n",
" freq_type=freq_type)\n",
"\n",
" return train_dataset, val_dataset\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Model Creation"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"def get_model(load_weights: bool = False):\n",
" device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
" repo_id = \"google/timesfm-2.0-500m-pytorch\"\n",
" hparams = TimesFmHparams(\n",
" backend=device,\n",
" per_core_batch_size=32,\n",
" horizon_len=128,\n",
" num_layers=50,\n",
" use_positional_embedding=False,\n",
" context_len=\n",
" 192, # Context length can be anything up to 2048 in multiples of 32\n",
" )\n",
" tfm = TimesFm(hparams=hparams,\n",
" checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))\n",
"\n",
" model = PatchedTimeSeriesDecoder(tfm._model_config)\n",
" if load_weights:\n",
" checkpoint_path = path.join(snapshot_download(repo_id), \"torch_model.ckpt\")\n",
" loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)\n",
" model.load_state_dict(loaded_checkpoint)\n",
" return model, hparams, tfm._model_config\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def plot_predictions(\n",
" model: TimesFm,\n",
" val_dataset: Dataset,\n",
" save_path: Optional[str] = \"predictions.png\",\n",
") -> None:\n",
" \"\"\"\n",
" Plot model predictions against ground truth for a batch of validation data.\n",
"\n",
" Args:\n",
" model: Trained TimesFM model\n",
" val_dataset: Validation dataset\n",
" save_path: Path to save the plot\n",
" \"\"\"\n",
" import matplotlib.pyplot as plt\n",
"\n",
" model.eval()\n",
"\n",
" x_context, x_padding, freq, x_future = val_dataset[0]\n",
" x_context = x_context.unsqueeze(0) # Add batch dimension\n",
" x_padding = x_padding.unsqueeze(0)\n",
" freq = freq.unsqueeze(0)\n",
" x_future = x_future.unsqueeze(0)\n",
"\n",
" device = next(model.parameters()).device\n",
" x_context = x_context.to(device)\n",
" x_padding = x_padding.to(device)\n",
" freq = freq.to(device)\n",
" x_future = x_future.to(device)\n",
"\n",
" with torch.no_grad():\n",
" predictions = model(x_context, x_padding.float(), freq)\n",
" predictions_mean = predictions[..., 0] # [B, N, horizon_len]\n",
" last_patch_pred = predictions_mean[:, -1, :] # [B, horizon_len]\n",
"\n",
" context_vals = x_context[0].cpu().numpy()\n",
" future_vals = x_future[0].cpu().numpy()\n",
" pred_vals = last_patch_pred[0].cpu().numpy()\n",
"\n",
" context_len = len(context_vals)\n",
" horizon_len = len(future_vals)\n",
"\n",
" plt.figure(figsize=(12, 6))\n",
"\n",
" plt.plot(range(context_len),\n",
" context_vals,\n",
" label=\"Historical Data\",\n",
" color=\"blue\",\n",
" linewidth=2)\n",
"\n",
" plt.plot(\n",
" range(context_len, context_len + horizon_len),\n",
" future_vals,\n",
" label=\"Ground Truth\",\n",
" color=\"green\",\n",
" linestyle=\"--\",\n",
" linewidth=2,\n",
" )\n",
"\n",
" plt.plot(range(context_len, context_len + horizon_len),\n",
" pred_vals,\n",
" label=\"Prediction\",\n",
" color=\"red\",\n",
" linewidth=2)\n",
"\n",
" plt.xlabel(\"Time Step\")\n",
" plt.ylabel(\"Value\")\n",
" plt.title(\"TimesFM Predictions vs Ground Truth\")\n",
" plt.legend()\n",
" plt.grid(True)\n",
"\n",
" if save_path:\n",
" plt.savefig(save_path)\n",
" print(f\"Plot saved to {save_path}\")\n",
"\n",
" plt.close()\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"def get_data(context_len: int,\n",
" horizon_len: int,\n",
" freq_type: int = 0) -> Tuple[Dataset, Dataset]:\n",
" df = yf.download(\"AAPL\", start=\"2010-01-01\", end=\"2019-01-01\")\n",
" time_series = df[\"Close\"].values\n",
"\n",
" train_dataset, val_dataset = prepare_datasets(\n",
" series=time_series,\n",
" context_length=context_len,\n",
" horizon_length=horizon_len,\n",
" freq_type=freq_type,\n",
" train_split=0.8,\n",
" )\n",
"\n",
" print(f\"Created datasets:\")\n",
" print(f\"- Training samples: {len(train_dataset)}\")\n",
" print(f\"- Validation samples: {len(val_dataset)}\")\n",
" print(f\"- Using frequency type: {freq_type}\")\n",
" return train_dataset, val_dataset\n",
"\n",
"\n",
"\n",
"def single_gpu_example():\n",
" \"\"\"Basic example of finetuning TimesFM on stock data.\"\"\"\n",
" model, hparams, tfm_config = get_model(load_weights=True)\n",
" config = FinetuningConfig(batch_size=256,\n",
" num_epochs=5,\n",
" learning_rate=1e-4,\n",
" use_wandb=True,\n",
" freq_type=1,\n",
" log_every_n_steps=10,\n",
" val_check_interval=0.5,\n",
" use_quantile_loss=True)\n",
"\n",
" train_dataset, val_dataset = get_data(128,\n",
" tfm_config.horizon_len,\n",
" freq_type=config.freq_type)\n",
" finetuner = TimesFMFinetuner(model, config)\n",
"\n",
" print(\"\\nStarting finetuning...\")\n",
" results = finetuner.finetune(train_dataset=train_dataset,\n",
" val_dataset=val_dataset)\n",
"\n",
" print(\"\\nFinetuning completed!\")\n",
" print(f\"Training history: {len(results['history']['train_loss'])} epochs\")\n",
"\n",
" plot_predictions(\n",
" model=model,\n",
" val_dataset=val_dataset,\n",
" save_path=\"timesfm_predictions.png\",\n",
" )\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ac84aeda3a1749ae8f30b06859067bb1",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Fetching 3 files: 0%| | 0/3 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "6d9d8081fc514c6d8601a2e0e63954a2",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Fetching 3 files: 0%| | 0/3 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"[*********************100%***********************] 1 of 1 completed\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Created datasets:\n",
"- Training samples: 1556\n",
"- Validation samples: 198\n",
"- Using frequency type: 1\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend. Please refer to https://wandb.me/wandb-core for more information.\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mmishacamry\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
]
},
{
"data": {
"text/html": [
"Tracking run with wandb version 0.19.1"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"Run data is saved locally in <code>/home/chertushkin/forks/timesfm/notebooks/wandb/run-20250217_114343-tjs63ml2</code>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"Syncing run <strong><a href='https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2' target=\"_blank\">chocolate-eon-50</a></strong> to <a href='https://wandb.ai/mishacamry/timesfm-finetuning' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
" View project at <a href='https://wandb.ai/mishacamry/timesfm-finetuning' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning</a>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
" View run at <a href='https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2</a>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Starting finetuning...\n"
]
},
{
"data": {
"text/html": [],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<br> <style><br> .wandb-row {<br> display: flex;<br> flex-direction: row;<br> flex-wrap: wrap;<br> justify-content: flex-start;<br> width: 100%;<br> }<br> .wandb-col {<br> display: flex;<br> flex-direction: column;<br> flex-basis: 100%;<br> flex: 1;<br> padding: 10px;<br> }<br> </style><br><div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>▁▃▅▆█</td></tr><tr><td>learning_rate</td><td>▁▁▁▁▁</td></tr><tr><td>train_loss</td><td>█▃▂▁▁</td></tr><tr><td>val_loss</td><td>█▁▄▁▂</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>5</td></tr><tr><td>learning_rate</td><td>0.0001</td></tr><tr><td>train_loss</td><td>2.85423</td></tr><tr><td>val_loss</td><td>26.7628</td></tr></table><br/></div></div>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
" View run <strong style=\"color:#cdcd00\">chocolate-eon-50</strong> at: <a href='https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning/runs/tjs63ml2</a><br> View project at: <a href='https://wandb.ai/mishacamry/timesfm-finetuning' target=\"_blank\">https://wandb.ai/mishacamry/timesfm-finetuning</a><br>Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"Find logs at: <code>./wandb/run-20250217_114343-tjs63ml2/logs</code>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Finetuning completed!\n",
"Training history: 5 epochs\n",
"Plot saved to timesfm_predictions.png\n"
]
}
],
"source": [
"single_gpu_example()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "timesfm-DnAbSweh-py3.11",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.10"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
v1/peft/README.md
+42
View File
@@ -0,0 +1,42 @@
# Fine-Tuning Pipeline
This folder contains a generic fine-tuning pipeline designed to support multiple PEFT fine-tuning strategies.
## Features
- **Supported Fine-Tuning Strategies**:
- **Full Fine-Tuning**: Adjusts all parameters of the model during training.
- **[Linear Probing](https://arxiv.org/abs/2302.11939)**: Fine-tunes only the residual blocks and the embedding layer, leaving other parameters unchanged.
- **[LoRA (Low-Rank Adaptation)](https://arxiv.org/abs/2106.09685)**: A memory-efficient method that fine-tunes a small number of parameters by decomposing the weight matrices into low-rank matrices.
- **[DoRA (Directional LoRA)](https://arxiv.org/abs/2402.09353v4)**: An extension of LoRA that decomposes pre-trained weights into magnitude and direction components. It uses LoRA for directional adaptation, enhancing learning capacity and stability without additional inference overhead.
## Usage
### Fine-Tuning Script
The provided finetune.py script allows you to fine-tune a model with specific configurations. You can customize various parameters to suit your dataset and desired fine-tuning strategy.
Example Usage:
```zsh
source finetune.sh
```
This script runs the finetune.py file with a predefined set of hyperparameters for the model. You can adjust the parameters in the script as needed.
### Available Options
Run the script with the --help flag to see a full list of available options and their descriptions:
```zsh
python3 finetune.py --help
```
Script Configuration
You can modify the following key parameters directly in the finetune.sh script:
Fine-Tuning Strategy: Toggle between full fine-tuning, LoRA \[`--use-lora`\], DoRA [\[`--use-dora`\]], or Linear Probing \[`--use-linear-probing`\].
### Performance Comparison
The figure below compares the performance of LoRA/DoRA against Linear Probing under the following conditions:
<img width="528" alt="image" src="https://github.com/user-attachments/assets/6c9f820b-5865-4821-8014-c346b9d632a5">
- Training data split: 60% train, 20% validation, 20% test.
- Benchmark: context_len=128, horizon_len=96
- Fine-tuning: context_len=128, horizon_len=128
- Black: Best result.
- Blue: Second best result.
v1/peft/finetune.py
+402
View File
@@ -0,0 +1,402 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Finetune pipeline.
"""
import gc
import logging
import warnings
from datetime import datetime
from typing import Tuple
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
import typer
import wandb
from jax import numpy as jnp
from paxml import checkpoint_types, checkpoints, learners, tasks_lib, trainer_lib
from praxis import optimizers, pax_fiddle, py_utils, schedules
from rich import print
from tqdm import tqdm
from typing_extensions import Annotated
from adapter.utils import get_adapter_params, load_adapter_layer
from timesfm import TimesFm, data_loader, patched_decoder
NestedMap = py_utils.NestedMap
warnings.filterwarnings("ignore")
cmdstanpy_logger = logging.getLogger("cmdstanpy")
absl_logger = logging.getLogger("absl")
cmdstanpy_logger.disabled = True
absl_logger.disabled = True
"""
TimesFM model config. These are fixed since pre-training was done
with this configuration.
"""
INPUT_PATCH_LEN = 32
OUTPUT_PATCH_LEN = 128
NUM_LAYERS = 20
MODEL_DIMS = 1280
QUANTILES = list(np.arange(1, 10) / 10.0)
EPS = 1e-7
RANDOM_SEED = 1234
def finetune(
*,
model_name: Annotated[
str, typer.Option(help="Specify the name of the huggingface model.")
] = "google/timesfm-1.0-200m",
checkpoint_path: Annotated[
str, typer.Option(help="The path to the local model checkpoint.")
] = None,
datetime_col: Annotated[str, typer.Option(help="Column having datetime.")] = "ds",
ts_cols: Annotated[
list[str], typer.Option(help="Columns of time-series features.")
] = [],
normalize: Annotated[
bool, typer.Option(help="Normalize data for eval or not")
] = True,
context_len: Annotated[int, typer.Option(help="Length of the context window")],
horizon_len: Annotated[int, typer.Option(help="Prediction length.")],
freq: Annotated[
str,
typer.Option(
...,
help="Frequency Map Str",
),
],
data_path: Annotated[str, typer.Option(help="Path to dataset csv")],
boundaries: Annotated[
Tuple[int, int, int],
typer.Option(
help="boundaries of dataset to train, val, test",
),
] = (0, 0, 0),
backend: Annotated[str, typer.Option(help="Backend device: cpu, gpu, tpu")],
batch_size: Annotated[
int, typer.Option(help="Batch size for the randomly sampled batch")
] = 16,
num_epochs: Annotated[int, typer.Option(help="Number of epochs")],
learning_rate: Annotated[float, typer.Option(help="adam optimizer learning rate")],
adam_epsilon: Annotated[float, typer.Option(help="adam optimizer epsilon")],
adam_clip_threshold: Annotated[
float, typer.Option(help="adam optimizer clip threshold")
],
cos_initial_decay_value: Annotated[
float, typer.Option(help="cosine initial decay value")
],
cos_final_decay_value: Annotated[
float, typer.Option(help="cosine final decay value")
],
cos_decay_steps: Annotated[int, typer.Option(help="Number of cosine decay steps")],
ema_decay: Annotated[float, typer.Option(help="Exponential moving average decay")],
early_stop_patience: Annotated[
int, typer.Option(..., help="Early stopping patience")
] = 5,
use_lora: Annotated[
bool,
typer.Option(
help="Train low rank adapters for stacked transformer block",
),
] = False,
lora_rank: Annotated[
int,
typer.Option(
help="LoRA Rank",
),
] = 8,
lora_target_modules: Annotated[
str,
typer.Option(
help="LoRA target modules of the transformer block. Allowed values: [all, attention, mlp]"
),
] = "all",
use_dora: Annotated[
bool,
typer.Option(
help="Apply DoRA strategy along with LoRA.",
),
] = False,
use_linear_probing: Annotated[
bool,
typer.Option(
help="Linear Probing. Train only input/output and embedding params. Freeze params in stack transformer block.",
),
] = False,
checkpoint_dir: Annotated[
str, typer.Option(help="Checkpoint directory")
] = "./checkpoints",
wandb_project: Annotated[
str, typer.Option(help="Weights & Biases project name")
] = "google_timesfm_finetune",
) -> None:
key = jax.random.PRNGKey(seed=RANDOM_SEED)
wandb.init(project=wandb_project, config=locals())
data_df = pd.read_csv(open(data_path, "r"))
if boundaries == (0, 0, 0):
# Default boundaries: train 60%, val 20%, test 20%
boundaries = [
int(len(data_df) * 0.6),
int(len(data_df) * 0.8),
len(data_df) - 1,
]
ts_cols = [col for col in data_df.columns if col != datetime_col]
dtl = data_loader.TimeSeriesdata(
data_path=data_path,
datetime_col=datetime_col,
num_cov_cols=None,
cat_cov_cols=None,
ts_cols=np.array(ts_cols),
train_range=[0, boundaries[0]],
val_range=[boundaries[0], boundaries[1]],
test_range=[boundaries[1], boundaries[2]],
hist_len=context_len,
pred_len=horizon_len,
batch_size=batch_size,
freq=freq,
normalize=normalize,
epoch_len=None,
holiday=False,
permute=False,
)
train_batches = dtl.tf_dataset(mode="train", shift=1).batch(batch_size)
val_batches = dtl.tf_dataset(mode="val", shift=horizon_len)
for tbatch in tqdm(train_batches.as_numpy_iterator()):
pass
tfm = TimesFm(
context_len=context_len,
horizon_len=horizon_len,
input_patch_len=INPUT_PATCH_LEN,
output_patch_len=OUTPUT_PATCH_LEN,
num_layers=NUM_LAYERS,
model_dims=MODEL_DIMS,
backend=backend,
per_core_batch_size=batch_size,
quantiles=QUANTILES,
)
if checkpoint_path:
tfm.load_from_checkpoint(
checkpoint_path=checkpoint_path,
checkpoint_type=checkpoints.CheckpointType.FLAX,
)
else:
tfm.load_from_checkpoint(
repo_id=model_name,
checkpoint_type=checkpoints.CheckpointType.FLAX,
)
model = pax_fiddle.Config(
patched_decoder.PatchedDecoderFinetuneModel,
name="patched_decoder_finetune",
core_layer_tpl=tfm.model_p,
)
if use_lora:
load_adapter_layer(
mdl_vars=tfm._train_state.mdl_vars,
model=model.core_layer_tpl,
lora_rank=lora_rank,
lora_target_modules=lora_target_modules,
use_dora=use_dora,
)
@pax_fiddle.auto_config
def build_learner() -> learners.Learner:
bprop_variable_inclusion = []
bprop_variable_exclusion = []
if use_lora:
bprop_variable_inclusion.append(r"^.*lora.*$")
if use_dora:
bprop_variable_inclusion.append(r"^.*dora.*$")
elif use_linear_probing:
bprop_variable_exclusion = [".*/stacked_transformer_layer/.*"]
return pax_fiddle.Config(
learners.Learner,
name="learner",
loss_name="avg_qloss",
optimizer=optimizers.Adam(
epsilon=adam_epsilon,
clip_threshold=adam_clip_threshold,
learning_rate=learning_rate,
lr_schedule=pax_fiddle.Config(
schedules.Cosine,
initial_value=cos_initial_decay_value,
final_value=cos_final_decay_value,
total_steps=cos_decay_steps,
),
ema_decay=ema_decay,
),
bprop_variable_exclusion=bprop_variable_exclusion,
bprop_variable_inclusion=bprop_variable_inclusion,
)
task_p = tasks_lib.SingleTask(
name="ts-learn",
model=model,
train=tasks_lib.SingleTask.Train(
learner=build_learner(),
),
)
task_p.model.ici_mesh_shape = [1, 1, 1]
task_p.model.mesh_axis_names = ["replica", "data", "mdl"]
DEVICES = np.array(jax.devices()).reshape([1, 1, 1])
jax.sharding.Mesh(DEVICES, ["replica", "data", "mdl"])
num_devices = jax.local_device_count()
print(f"num_devices: {num_devices}")
print(f"device kind: {jax.local_devices()[0].device_kind}")
jax_task = task_p
key, init_key = jax.random.split(key)
def process_train_batch(batch):
past_ts = batch[0].reshape(batch_size * len(ts_cols), -1)
actual_ts = batch[3].reshape(batch_size * len(ts_cols), -1)
return NestedMap(input_ts=past_ts, actual_ts=actual_ts)
def process_eval_batch(batch):
past_ts = batch[0]
actual_ts = batch[3]
return NestedMap(input_ts=past_ts, actual_ts=actual_ts)
jax_model_states, _ = trainer_lib.initialize_model_state(
jax_task,
init_key,
process_train_batch(tbatch),
checkpoint_type=checkpoint_types.CheckpointType.GDA,
)
jax_model_states.mdl_vars["params"]["core_layer"] = tfm._train_state.mdl_vars[
"params"
]
gc.collect()
jax_task = task_p
def train_step(states, prng_key, inputs):
return trainer_lib.train_step_single_learner(jax_task, states, prng_key, inputs)
def eval_step(states, prng_key, inputs):
states = states.to_eval_state()
return trainer_lib.eval_step_single_learner(jax_task, states, prng_key, inputs)
key, train_key, eval_key = jax.random.split(key, 3)
train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())
eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())
p_train_step = jax.pmap(train_step, axis_name="batch")
p_eval_step = jax.pmap(eval_step, axis_name="batch")
replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)
def reshape_batch_for_pmap(batch, num_devices):
def _reshape(input_tensor):
bsize = input_tensor.shape[0]
residual_shape = list(input_tensor.shape[1:])
nbsize = bsize // num_devices
return jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)
return jax.tree.map(_reshape, batch)
patience = 0
best_eval_loss = 1e7
checkpoint_dir = f"{checkpoint_dir}/run_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{wandb.run.id}"
for epoch in range(num_epochs):
if patience >= early_stop_patience:
print("Early stopping.")
break
print(f"Epoch: {epoch + 1}")
train_its = train_batches.as_numpy_iterator()
train_losses = []
for batch in tqdm(train_its):
tbatch = process_train_batch(batch)
tbatch = reshape_batch_for_pmap(tbatch, num_devices)
replicated_jax_states, step_fun_out = p_train_step(
replicated_jax_states, train_prng_seed, tbatch
)
train_losses.append(step_fun_out.loss[0])
wandb.log({"train_step_loss": step_fun_out.loss[0]})
avg_train_loss = np.mean(train_losses)
print("Starting eval.")
val_its = val_batches.as_numpy_iterator()
eval_losses = []
for ev_batch in tqdm(val_its):
ebatch = process_eval_batch(ev_batch)
ebatch = reshape_batch_for_pmap(ebatch, num_devices)
_, step_fun_out = p_eval_step(replicated_jax_states, eval_prng_seed, ebatch)
eval_losses.append(step_fun_out.loss[0])
wandb.log({"eval_step_loss": step_fun_out.loss[0]})
avg_eval_loss = np.mean(eval_losses)
print(f"Train Loss: {avg_train_loss}, Val Loss: {avg_eval_loss}")
wandb.log(
{
"epoch": epoch + 1,
"avg_train_loss": avg_train_loss,
"avg_val_loss": avg_eval_loss,
}
)
if avg_eval_loss < best_eval_loss or np.isnan(avg_eval_loss):
best_eval_loss = avg_eval_loss
print("Saving checkpoint.")
jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(
replicated_jax_states
)
if use_lora:
adapter_params = get_adapter_params(
params=jax_state_for_saving.mdl_vars,
lora_target_modules=lora_target_modules,
num_layers=NUM_LAYERS,
use_dora=use_dora,
)
jax_state_for_saving.mdl_vars["params"] = adapter_params
checkpoints.save_checkpoint(
jax_state_for_saving, checkpoint_dir, overwrite=True
)
patience = 0
del jax_state_for_saving
gc.collect()
else:
patience += 1
print(f"patience: {patience}")
print("Fine-tuning completed.")
if __name__ == "__main__":
typer.run(finetune)
v1/peft/finetune.sh
+31
View File
@@ -0,0 +1,31 @@
#!/bin/bash
# Script to finetune a model with specific configurations
# Adjust the parameters below as needed. For a full list of options and descriptions, run the script with the --help flag.
export TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false
python3 finetune.py \
--model-name="google/timesfm-1.0-200m" \
--backend="cpu" \
--horizon-len=128 \
--context-len=512 \
--freq="15min" \
--data-path="../datasets/ETT-small/ETTm1.csv" \
--num-epochs=100 \
--learning-rate=1e-3 \
--adam-epsilon=1e-7 \
--adam-clip-threshold=1e2 \
--early-stop-patience=10 \
--datetime-col="date" \
--use-lora \
--lora-rank=1 \
--lora-target-modules="all" \
--use-dora \
--cos-initial-decay-value=1e-4 \
--cos-decay-steps=40000 \
--cos-final-decay-value=1e-5 \
--ema-decay=0.9999
# To see all available options and their descriptions, use the --help flag
# python3 finetune.py --help
v1/peft/usage.ipynb
+203
View File
@@ -0,0 +1,203 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load Base Model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from timesfm import TimesFm, freq_map, data_loader\n",
"from adapter.utils import load_adapter_checkpoint\n",
"from tqdm import tqdm\n",
"import numpy as np\n",
"import pandas as pd\n",
"\n",
"\n",
"tfm = TimesFm(\n",
" context_len=512,\n",
" horizon_len=128,\n",
" input_patch_len=32,\n",
" output_patch_len=128,\n",
" num_layers=20,\n",
" model_dims=1280,\n",
" backend=\"cpu\",\n",
")\n",
"tfm.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"DATA_DICT = {\n",
" \"ettm2\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"ettm1\": {\n",
" \"boundaries\": [34560, 46080, 57600],\n",
" \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n",
" \"freq\": \"15min\",\n",
" },\n",
" \"etth2\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"etth1\": {\n",
" \"boundaries\": [8640, 11520, 14400],\n",
" \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"elec\": {\n",
" \"boundaries\": [18413, 21044, 26304],\n",
" \"data_path\": \"../datasets/electricity/electricity.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"traffic\": {\n",
" \"boundaries\": [12280, 14036, 17544],\n",
" \"data_path\": \"../datasets/traffic/traffic.csv\",\n",
" \"freq\": \"H\",\n",
" },\n",
" \"weather\": {\n",
" \"boundaries\": [36887, 42157, 52696],\n",
" \"data_path\": \"../datasets/weather/weather.csv\",\n",
" \"freq\": \"10min\",\n",
" },\n",
"}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load Adapter Checkpoint\n",
"\n",
"Specify the adapter checkpoint path, rank and the modules used to train the adapters and whether dora was employed or not."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"load_adapter_checkpoint(\n",
" model=tfm,\n",
" adapter_checkpoint_path=\"./checkpoints/run_20240716_163900_lyo4psz3\",\n",
" lora_rank=1,\n",
" lora_target_modules=\"all\",\n",
" use_dora=True,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Test Performance"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dataset = \"ettm1\"\n",
"data_path = DATA_DICT[dataset][\"data_path\"]\n",
"freq = DATA_DICT[dataset][\"freq\"]\n",
"int_freq = freq_map(freq)\n",
"boundaries = DATA_DICT[dataset][\"boundaries\"]\n",
"\n",
"data_df = pd.read_csv(open(data_path, \"r\"))\n",
"\n",
"ts_cols = [col for col in data_df.columns if col != \"date\"]\n",
"num_cov_cols = None\n",
"cat_cov_cols = None\n",
"\n",
"context_len = 512\n",
"pred_len = 96\n",
"\n",
"num_ts = len(ts_cols)\n",
"batch_size = 16\n",
"\n",
"dtl = data_loader.TimeSeriesdata(\n",
" data_path=data_path,\n",
" datetime_col=\"date\",\n",
" num_cov_cols=num_cov_cols,\n",
" cat_cov_cols=cat_cov_cols,\n",
" ts_cols=np.array(ts_cols),\n",
" train_range=[0, boundaries[0]],\n",
" val_range=[boundaries[0], boundaries[1]],\n",
" test_range=[boundaries[1], boundaries[2]],\n",
" hist_len=context_len,\n",
" pred_len=pred_len,\n",
" batch_size=num_ts,\n",
" freq=\"15min\",\n",
" normalize=True,\n",
" epoch_len=None,\n",
" holiday=False,\n",
" permute=True,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mae_losses = []\n",
"for batch in tqdm(test_batches.as_numpy_iterator()):\n",
" past = batch[0]\n",
" actuals = batch[3]\n",
" _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n",
" forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n",
" mae_losses.append(np.abs(forecasts - actuals).mean())\n",
"\n",
"print(f\"MAE: {np.mean(mae_losses)}\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "tanmay_tfm_env",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
v1/poetry.lock
Generated
+7282
View File
File diff suppressed because it is too large Load Diff
v1/pyproject.toml
+75
View File
@@ -0,0 +1,75 @@
[tool.poetry]
name = "timesfm"
packages = [
{ include = "timesfm", from = "src" },
{ include = "finetuning", from = "src" },
]
description = "Open weights time-series foundation model from Google Research."
version = "1.3.0"
authors = [
"Rajat Sen <senrajat@google.com>",
"Yichen Zhou <yichenzhou@google.com>",
"Abhimanyu Das <abhidas@google.com>",
"Petros Mol <pmol@google.com>",
"Justin Güse <guese.justin@gmail.com>",
"Michael Chertushkin <chertushkinmichael@gmail.com>"
]
readme = "README.md"
keywords = ["time series", "timesfm", "forecast", "time series model"]
homepage = "https://github.com/google-research/timesfm"
repository = "https://github.com/google-research/timesfm"
classifiers = [
"Environment :: Console",
"Framework :: Flake8",
"Operating System :: OS Independent",
"Topic :: Software Development :: Documentation",
"Topic :: Software Development :: Libraries :: Python Modules",
"Topic :: Software Development :: Quality Assurance",
]
include = ["LICENSE"]
[tool.poetry.dependencies]
python = ">=3.10,<3.12"
einshape = ">=1.0.0"
numpy = ">=1.26.4"
pandas = ">=2.0.0"
utilsforecast = ">=0.1.10"
huggingface_hub = { version = ">=0.23.0", extras = ["cli"] }
scikit-learn = ">=1.2.2"
typer = ">=0.12.3"
wandb = ">=0.17.5"
absl-py = ">=1.4.0"
safetensors = "^0.5.3"
[tool.poetry.extras]
pax = ["paxml", "lingvo", "jax", "jaxlib"]
torch = ["torch"]
[tool.poetry.dependencies.paxml]
version = ">=1.4.0"
python = ">=3.10,<3.11"
[tool.poetry.dependencies.lingvo]
version = ">=0.12.7"
python = ">=3.10,<3.11"
[tool.poetry.dependencies.jax]
version = ">=0.4.26"
extras = ["cuda12"]
python = ">=3.10,<3.12" # Support both python versions
[tool.poetry.dependencies.jaxlib]
version = ">=0.4.26"
python = ">=3.10,<3.12" # Support both python versions
[tool.poetry.dependencies.torch]
version = ">=2.0.0"
extras = ["cuda"]
python = ">=3.11,<3.12"
[tool.poetry.group.dev.dependencies]
pytest = ">=8.3.2"
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
v1/src/adapter/__init__.py
+18
View File
@@ -0,0 +1,18 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""adapter init file."""
from .dora_layers import DoraAttentionProjection, DoraCombinedQKVProjection, DoraLinear
from .lora_layers import LoraAttentionProjection, LoraCombinedQKVProjection, LoraLinear
v1/src/adapter/dora_layers.py
+202
View File
@@ -0,0 +1,202 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from jax import numpy as jnp
from praxis import base_layer
from praxis.layers import attentions, linears
WeightInit = base_layer.WeightInit
WeightHParams = base_layer.WeightHParams
class DoraTheta(base_layer.Theta):
def __init__(self, module):
self.module = module
def _dora_initialized(self):
if (
self.module.has_variable("params", "lora_a")
and self.module.has_variable("params", "lora_b")
and self.module.has_variable("params", "dora_m")
and "lora_a" in self.module._weight_hparams
and "lora_b" in self.module._weight_hparams
and "dora_m" in self.module._weight_hparams
):
return True
else:
return False
def _dorafy_var(self, w):
lora_a = super().__getattr__("lora_a")
lora_b = super().__getattr__("lora_b")
dora_m = super().__getattr__("dora_m")
lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
lora_delta = jnp.reshape(lora_delta, w.shape)
w_prime = w + lora_delta
column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
norm_adapted = w_prime / column_norm
w_prime = dora_m * norm_adapted
return w_prime
def __getattr__(self, k):
var = super().__getattr__(k)
if not self._dora_initialized():
return var
if k == "w":
return self._dorafy_var(var)
return var
def __getitem__(self, k):
var = super().__getattr__(k)
if not self._dora_initialized():
return var
if k == "w":
return self._dorafy_var(var)
return var
class DoraThetaDescriptor:
"""Dot syntax accession descriptor."""
def __get__(self, obj, objtype=None):
return DoraTheta(obj)
class DoraLinear(linears.Linear):
rank: int = 0
lora_init: WeightInit | None = None
theta = DoraThetaDescriptor()
def setup(self) -> None:
lora_init = self.lora_init if self.lora_init else self.weight_init
super().setup()
self.create_variable(
"lora_a",
WeightHParams(
shape=[self.input_dims, self.rank],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None],
),
)
self.create_variable(
"lora_b",
WeightHParams(
shape=[self.output_dims, self.rank],
init=WeightInit.Constant(scale=0.0),
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None],
),
)
self.create_variable(
"dora_m",
WeightHParams(
shape=[1, self.output_dims],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None],
),
)
class DoraAttentionProjection(attentions.AttentionProjection):
rank: int = 0
lora_init: WeightInit | None = None
theta = DoraThetaDescriptor()
def setup(self) -> None:
super().setup()
w_weight_params = self._weight_hparams["w"]
lora_init = self.lora_init if self.lora_init else w_weight_params.init
self.create_variable(
"lora_a",
WeightHParams(
shape=[self.input_dim, self.rank],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[
None,
None,
],
),
)
self.create_variable(
"lora_b",
WeightHParams(
shape=[self.dim_per_head * self.num_heads, self.rank],
init=WeightInit.Constant(scale=0.0),
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[
None,
None,
],
),
)
self.create_variable(
"dora_m",
WeightHParams(
shape=[1, self.num_heads, self.dim_per_head],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None, None],
),
)
class DoraCombinedQKVProjection(attentions.CombinedQKVProjectionLayer):
rank: int = 0
lora_init: WeightInit | None = None
theta = DoraThetaDescriptor()
def setup(self) -> None:
super().setup()
w_weight_params = self._weight_hparams["w"]
lora_init = self.lora_init if self.lora_init else w_weight_params.init
self.create_variable(
"lora_a",
WeightHParams(
shape=[3, self.input_dim, self.rank],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None, None],
),
)
self.create_variable(
"lora_b",
WeightHParams(
shape=[3, self.dim_per_head * self.num_heads, self.rank],
init=WeightInit.Constant(scale=0.0),
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None, None],
),
)
self.create_variable(
"dora_m",
WeightHParams(
shape=[3, 1, self.num_heads, self.dim_per_head],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None, None, None],
),
)
v1/src/adapter/lora_layers.py
+166
View File
@@ -0,0 +1,166 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from jax import numpy as jnp
from praxis import base_layer
from praxis.layers import attentions, linears
WeightInit = base_layer.WeightInit
WeightHParams = base_layer.WeightHParams
class LoraTheta(base_layer.Theta):
def __init__(self, module):
self.module = module
def _lora_initialized(self):
if (
self.module.has_variable("params", "lora_a")
and self.module.has_variable("params", "lora_b")
and "lora_a" in self.module._weight_hparams
and "lora_b" in self.module._weight_hparams
):
return True
else:
return False
def _lorafy_var(self, w):
lora_a = super().__getattr__("lora_a")
lora_b = super().__getattr__("lora_b")
lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
lora_delta = jnp.reshape(lora_delta, w.shape)
w_prime = w + lora_delta
return w_prime
def __getattr__(self, k):
var = super().__getattr__(k)
if not self._lora_initialized():
return var
if k == "w":
return self._lorafy_var(var)
return var
def __getitem__(self, k):
var = super().__getattr__(k)
if not self._lora_initialized():
return var
if k == "w":
return self._lorafy_var(var)
return var
class LoraThetaDescriptor:
"""Dot syntax accession descriptor."""
def __get__(self, obj, objtype=None):
return LoraTheta(obj)
class LoraLinear(linears.Linear):
rank: int = 0
lora_init: WeightInit | None = None
theta = LoraThetaDescriptor()
def setup(self) -> None:
lora_init = self.lora_init if self.lora_init else self.weight_init
super().setup()
self.create_variable(
"lora_a",
WeightHParams(
shape=[self.input_dims, self.rank],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None],
),
)
self.create_variable(
"lora_b",
WeightHParams(
shape=[self.output_dims, self.rank],
init=WeightInit.Constant(scale=0.0),
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None],
),
)
class LoraAttentionProjection(attentions.AttentionProjection):
rank: int = 0
lora_init: WeightInit | None = None
theta = LoraThetaDescriptor()
def setup(self) -> None:
super().setup()
w_weight_params = self._weight_hparams["w"]
lora_init = self.lora_init if self.lora_init else w_weight_params.init
self.create_variable(
"lora_a",
WeightHParams(
shape=[self.input_dim, self.rank],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[
None,
None,
],
),
)
self.create_variable(
"lora_b",
WeightHParams(
shape=[self.dim_per_head * self.num_heads, self.rank],
init=WeightInit.Constant(scale=0.0),
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[
None,
None,
],
),
)
class LoraCombinedQKVProjection(attentions.CombinedQKVProjectionLayer):
rank: int = 0
lora_init: WeightInit | None = None
theta = LoraThetaDescriptor()
def setup(self) -> None:
super().setup()
w_weight_params = self._weight_hparams["w"]
lora_init = self.lora_init if self.lora_init else w_weight_params.init
self.create_variable(
"lora_a",
WeightHParams(
shape=[3, self.input_dim, self.rank],
init=lora_init,
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None, None],
),
)
self.create_variable(
"lora_b",
WeightHParams(
shape=[3, self.dim_per_head * self.num_heads, self.rank],
init=WeightInit.Constant(scale=0.0),
mesh_shape=self.mesh_shape,
tensor_split_dims_mapping=[None, None, None],
),
)
v1/src/adapter/utils.py
+487
View File
@@ -0,0 +1,487 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This file provides functionality for loading and merging adapter weights
in timesfm model, specifically for LoRA and DoRA.
LoRA: https://arxiv.org/abs/2106.09685
DoRA: https://arxiv.org/abs/2402.09353v4
"""
import time
import jax
import jax.numpy as jnp
from paxml import checkpoints, tasks_lib
from paxml.train_states import TrainState
from praxis import pax_fiddle
from adapter.dora_layers import (
DoraAttentionProjection,
DoraCombinedQKVProjection,
DoraLinear,
)
from adapter.lora_layers import (
LoraAttentionProjection,
LoraCombinedQKVProjection,
LoraLinear,
)
from timesfm import TimesFm
def get_adapter_params(
params: dict, lora_target_modules: str, num_layers: int, use_dora: bool = False
) -> dict:
"""
Extracts adapter parameters from the given model parameters for saving the checkpoint.
Args:
params (dict): The full model parameters.
lora_target_modules (str): Target modules for LoRA/DoRA adaptation.
num_layers (int): Number of transformer layers.
use_dora (bool, optional): Whether DoRA was used or not. Defaults to False.
Returns:
dict: A dictionary containing the extracted adapter parameters.
"""
adapter_params = {}
for i in range(num_layers):
layer_key = f"x_layers_{i}"
adapter_params[layer_key] = {}
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
linear = params["params"]["core_layer"]["stacked_transformer_layer"][
layer_key
]["ff_layer"][ff_layer_key]["linear"]
lora_a = linear["lora_a"]
lora_b = linear["lora_b"]
adapter_params[layer_key][ff_layer_key] = {
"lora_a": lora_a,
"lora_b": lora_b,
}
if use_dora:
adapter_params[layer_key][ff_layer_key]["dora_m"] = linear["dora_m"]
if lora_target_modules in ["all", "attention"]:
attention = params["params"]["core_layer"]["stacked_transformer_layer"][
layer_key
]["self_attention"]
for component in ["key", "query", "value", "post"]:
lora_a = attention[component]["lora_a"]
lora_b = attention[component]["lora_b"]
adapter_params[layer_key][component] = {
"lora_a": lora_a,
"lora_b": lora_b,
}
if use_dora:
adapter_params[layer_key][component]["dora_m"] = attention[
component
]["dora_m"]
return adapter_params
def load_adapter_checkpoint(
model: TimesFm,
adapter_checkpoint_path: str,
lora_rank: int,
lora_target_modules: str,
use_dora: bool,
) -> None:
"""
Loads an adapter checkpoint and merges it with the original model weights.
Args:
model (TimesFm): The model to update.
adapter_checkpoint_path (str): Path to the adapter checkpoint.
lora_rank (int): Rank of the LoRA adaptation.
lora_target_modules (str): Target modules for adaptation.
use_dora (bool): Whether DoRA was used or not.
Returns:
None
"""
"""
currently loading and initializing the model with adapter layers first and then merging the
adapter weights to original weights and replacing the adapter layers back to original layer.
# NOTE: refactor this. there should be a better way to load the LoRA checkpoint.
"""
model._logging(f"Restoring adapter checkpoint from {adapter_checkpoint_path}.")
start_time = time.time()
original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl = (
load_adapter_layer(
mdl_vars=model._train_state.mdl_vars,
model=model._model,
lora_rank=lora_rank,
lora_target_modules=lora_target_modules,
use_dora=use_dora,
)
)
var_weight_hparams = model._model.abstract_init_with_metadata(
model._get_sample_inputs(), do_eval=True
)
adapter_weight_hparams = _get_adapter_weight_params(
var_weight_hparams=var_weight_hparams,
lora_target_modules=lora_target_modules,
num_layers=model._model.stacked_transformer_params_tpl.num_layers,
use_dora=use_dora,
)
adapter_state_partition_specs = tasks_lib.create_state_partition_specs(
adapter_weight_hparams,
mesh_shape=model.mesh_shape,
mesh_axis_names=model.mesh_name,
discard_opt_states=True,
learners=None,
)
adapter_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
adapter_weight_hparams,
discard_opt_states=True,
learners=None,
)
adapter_train_state = checkpoints.restore_checkpoint(
state_global_shapes=adapter_state_local_shapes,
checkpoint_dir=adapter_checkpoint_path,
checkpoint_type=checkpoints.CheckpointType.FLAX,
state_specs=adapter_state_partition_specs,
step=None,
)
# add adapter weights to the original weights
_merge_adapter_weights(
model=model,
adapter_train_state=adapter_train_state,
lora_target_modules=lora_target_modules,
num_layers=model._model.stacked_transformer_params_tpl.num_layers,
use_dora=use_dora,
)
# replace back with the original model layer
if lora_target_modules in ["all", "mlp"]:
model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
original_linear_tpl
)
if lora_target_modules in ["all", "attention"]:
model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
original_attn_tpl
)
model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
original_combined_qkv_tpl
)
model._logging(
f"Restored adapter checkpoint in {time.time() - start_time:.2f} seconds."
)
# jit compile the model
model.jit_decode()
def _merge_adapter_weights(
model: TimesFm,
adapter_train_state: TrainState,
lora_target_modules: str,
num_layers: int,
use_dora: bool,
) -> None:
"""
Merges adapter weights with the original model weights.
Args:
model (TimesFm): The model to update.
adapter_train_state (TrainState): The adapter's train state.
lora_target_modules (str): Target modules for adaptation.
num_layers (int): Number of transformer layers.
use_dora (bool): Whether DoRA was used or not.
"""
for i in range(num_layers):
layer_key = f"x_layers_{i}"
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
linear = model._train_state.mdl_vars["params"][
"stacked_transformer_layer"
][layer_key]["ff_layer"][ff_layer_key]["linear"]
params = adapter_train_state.mdl_vars[layer_key][ff_layer_key]
lora_a = params["lora_a"]
lora_b = params["lora_b"]
w = linear["w"]
lora_delta = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
lora_delta = jnp.reshape(lora_delta, w.shape)
w_prime = w + lora_delta
if use_dora:
dora_m = params["dora_m"]
column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
norm_adapted = w_prime / column_norm
w_prime = dora_m * norm_adapted
linear["w"] = w_prime
del linear["dora_m"]
else:
linear["w"] = w_prime
del linear["lora_a"]
del linear["lora_b"]
if lora_target_modules in ["all", "attention"]:
attention = model._train_state.mdl_vars["params"][
"stacked_transformer_layer"
][layer_key]["self_attention"]
for component in ["key", "query", "value", "post"]:
params = adapter_train_state.mdl_vars[layer_key][component]
lora_a = params["lora_a"]
lora_b = params["lora_b"]
w = attention[component]["w"]
lora_delta = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
lora_delta = jnp.reshape(lora_delta, w.shape)
w_prime = w + lora_delta
if use_dora:
dora_m = params["dora_m"]
column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
norm_adapted = w_prime / column_norm
w_prime = dora_m * norm_adapted
attention[component]["w"] = w_prime
del attention[component]["dora_m"]
else:
attention[component]["w"] = w_prime
del attention[component]["lora_a"]
del attention[component]["lora_b"]
def _get_adapter_weight_params(
var_weight_hparams: dict, lora_target_modules: str, num_layers: int, use_dora: bool
) -> dict:
"""
Extracts adapter weight parameters from the given variable weight hyperparameters.
Args:
var_weight_hparams (dict): Variable weight hyperparameters.
lora_target_modules (str): Target modules for adaptation.
num_layers (int): Number of transformer layers.
use_dora (bool): Whether DoRA was used or not.
Returns:
dict: A dictionary containing the extracted adapter weight parameters.
"""
adapter_params = {}
for i in range(num_layers):
layer = f"x_layers_{i}"
adapter_params[layer] = {}
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
adapter_weight_params = var_weight_hparams["params"][
"stacked_transformer_layer"
][layer]["ff_layer"][ff_layer_key]["linear"]
adapter_params[layer][ff_layer_key] = {
"lora_a": adapter_weight_params["lora_a"],
"lora_b": adapter_weight_params["lora_b"],
}
if use_dora:
adapter_params[layer][ff_layer_key]["dora_m"] = (
adapter_weight_params["dora_m"]
)
if lora_target_modules in ["all", "attention"]:
for component in ["key", "value", "query", "post"]:
adapter_weight_params = var_weight_hparams["params"][
"stacked_transformer_layer"
][layer]["self_attention"][component]
adapter_params[layer][component] = {
"lora_a": adapter_weight_params["lora_a"],
"lora_b": adapter_weight_params["lora_b"],
}
if use_dora:
adapter_params[layer][component]["dora_m"] = adapter_weight_params[
"dora_m"
]
return adapter_params
def load_adapter_layer(
mdl_vars: dict,
model: pax_fiddle.Config,
lora_rank: int,
lora_target_modules: str,
use_dora: bool = False,
) -> tuple[pax_fiddle.Config, pax_fiddle.Config]:
"""
Updates target modules with adapter layers.
Args:
mdl_vars (dict): Model variables.
model (pax_fiddle.Config): Model configuration.
lora_rank (int): Rank of the LoRA adaptation.
lora_target_modules (str): Target modules for adaptation.
use_dora (bool, optional): Whether DoRA was used or not.
Returns:
tuple[pax_fiddle.Config, pax_fiddle.Config]: Updated model configurations.
"""
original_linear_tpl = original_attn_tpl = original_combined_qkv_tpl = None
if lora_target_modules in ["all", "mlp"]:
original_linear_tpl = (
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl
)
adapter_linear_tpl = (
pax_fiddle.Config(
DoraLinear,
rank=lora_rank,
)
if use_dora
else pax_fiddle.Config(
LoraLinear,
rank=lora_rank,
)
)
adapter_linear_tpl.copy_fields_from(original_linear_tpl)
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
adapter_linear_tpl
)
if lora_target_modules in ["all", "attention"]:
original_attn_tpl = (
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl
)
adapter_attn_tpl = (
pax_fiddle.Config(DoraAttentionProjection, rank=lora_rank)
if use_dora
else pax_fiddle.Config(LoraAttentionProjection, rank=lora_rank)
)
adapter_attn_tpl.copy_fields_from(original_attn_tpl)
original_combined_qkv_tpl = (
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl
)
adapter_combined_qkv_tpl = (
pax_fiddle.Config(DoraCombinedQKVProjection, rank=lora_rank)
if use_dora
else pax_fiddle.Config(LoraCombinedQKVProjection, rank=lora_rank)
)
adapter_combined_qkv_tpl.copy_fields_from(original_combined_qkv_tpl)
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
adapter_attn_tpl
)
model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
adapter_combined_qkv_tpl
)
# initialize and add adapter weights
_initialize_adapter_params(
mdl_vars=mdl_vars,
num_layers=model.stacked_transformer_params_tpl.num_layers,
lora_rank=lora_rank,
lora_target_modules=lora_target_modules,
use_dora=use_dora,
)
return original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl
def _initialize_adapter_params(
mdl_vars: dict,
num_layers,
lora_rank: int,
lora_target_modules: str,
use_dora: bool = False,
seed: int = 1234,
) -> dict:
"""
Initializes and adds adapter parameters to target modules.
Args:
mdl_vars (dict): Model variables.
num_layers (int): Number of transformer layers.
lora_rank (int): Rank of the LoRA adaptation.
lora_target_modules (str): Target modules for adaptation.
use_dora (bool, optional): Whether DoRA was used or not.
seed (int, optional): Random seed for initialization. Defaults to 1234.
Returns:
dict: Updated model variables with initialized adapter parameters.
"""
for i in range(num_layers):
layer_key = f"x_layers_{i}"
if lora_target_modules in ["all", "mlp"]:
for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
linear = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
"ff_layer"
][ff_layer_key]["linear"]
original_w = linear["w"]
input_dim, output_dim = original_w.shape
std_dev = 1 / jnp.sqrt(lora_rank)
normal_initializer = jax.nn.initializers.normal(std_dev)
lora_a = normal_initializer(
jax.random.key(seed), (input_dim, lora_rank), jnp.float32
)
lora_b = jnp.zeros((output_dim, lora_rank))
linear["lora_a"] = lora_a
linear["lora_b"] = lora_b
if use_dora:
norm = jnp.linalg.norm(original_w, ord=2, axis=0, keepdims=True)
linear["dora_m"] = norm
if lora_target_modules in ["all", "attention"]:
attention = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
"self_attention"
]
for component in ["key", "query", "value", "post"]:
original_w = attention[component]["w"]
w_dim = original_w.shape[0]
std_dev = 1 / jnp.sqrt(lora_rank)
normal_initializer = jax.nn.initializers.normal(std_dev)
lora_a = normal_initializer(
jax.random.key(seed), (w_dim, lora_rank), jnp.float32
)
lora_b = jnp.zeros((w_dim, lora_rank))
attention[component]["lora_a"] = lora_a
attention[component]["lora_b"] = lora_b
if use_dora:
norm = jnp.linalg.norm(
original_w, ord=2, axis=0, keepdims=True
).astype(jnp.float32)
attention[component]["dora_m"] = norm
return mdl_vars
v1/src/finetuning/__init__.py
View File
v1/src/finetuning/finetuning_example.py
+403
View File
@@ -0,0 +1,403 @@
"""
Example usage of the TimesFM Finetuning Framework.
For single GPU:
python script.py --training_mode=single
For multiple GPUs:
python script.py --training_mode=multi --gpu_ids=0,1,2
"""
import os
from dataclasses import asdict
from os import path
from typing import Optional, Tuple
import numpy as np
import pandas as pd
import torch
import torch.multiprocessing as mp
import yfinance as yf
from absl import app, flags
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from torch.utils.data import Dataset
from finetuning.finetuning_torch import FinetuningConfig, TimesFMFinetuner
from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams
from timesfm.pytorch_patched_decoder import (PatchedTimeSeriesDecoder,
TimesFMConfig)
FLAGS = flags.FLAGS
flags.DEFINE_enum(
"training_mode",
"single",
["single", "multi"],
'Training mode: "single" for single-GPU or "multi" for multi-GPU training.',
)
flags.DEFINE_list(
"gpu_ids", ["0"],
"Comma-separated list of GPU IDs to use for multi-GPU training. Example: 0,1,2"
)
flags.DEFINE_string(
"local_model_path",
None,
"Path to a local .safetensors model file. If provided, overrides Hugging Face download."
)
class TimeSeriesDataset(Dataset):
"""Dataset for time series data compatible with TimesFM."""
def __init__(self,
series: np.ndarray,
context_length: int,
horizon_length: int,
freq_type: int = 0):
"""
Initialize dataset.
Args:
series: Time series data
context_length: Number of past timesteps to use as input
horizon_length: Number of future timesteps to predict
freq_type: Frequency type (0, 1, or 2)
"""
if freq_type not in [0, 1, 2]:
raise ValueError("freq_type must be 0, 1, or 2")
self.series = series
self.context_length = context_length
self.horizon_length = horizon_length
self.freq_type = freq_type
self._prepare_samples()
def _prepare_samples(self) -> None:
"""Prepare sliding window samples from the time series."""
self.samples = []
total_length = self.context_length + self.horizon_length
for start_idx in range(0, len(self.series) - total_length + 1):
end_idx = start_idx + self.context_length
x_context = self.series[start_idx:end_idx]
x_future = self.series[end_idx:end_idx + self.horizon_length]
self.samples.append((x_context, x_future))
def __len__(self) -> int:
return len(self.samples)
def __getitem__(
self, index: int
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
x_context, x_future = self.samples[index]
x_context = torch.tensor(x_context, dtype=torch.float32)
x_future = torch.tensor(x_future, dtype=torch.float32)
input_padding = torch.zeros_like(x_context)
freq = torch.tensor([self.freq_type], dtype=torch.long)
return x_context, input_padding, freq, x_future
def prepare_datasets(series: np.ndarray,
context_length: int,
horizon_length: int,
freq_type: int = 0,
train_split: float = 0.8) -> Tuple[Dataset, Dataset]:
"""
Prepare training and validation datasets from time series data.
Args:
series: Input time series data
context_length: Number of past timesteps to use
horizon_length: Number of future timesteps to predict
freq_type: Frequency type (0, 1, or 2)
train_split: Fraction of data to use for training
Returns:
Tuple of (train_dataset, val_dataset)
"""
train_size = int(len(series) * train_split)
train_data = series[:train_size]
val_data = series[train_size:]
# Create datasets with specified frequency type
train_dataset = TimeSeriesDataset(train_data,
context_length=context_length,
horizon_length=horizon_length,
freq_type=freq_type)
val_dataset = TimeSeriesDataset(val_data,
context_length=context_length,
horizon_length=horizon_length,
freq_type=freq_type)
return train_dataset, val_dataset
def get_model(load_weights: bool = False):
device = "cuda" if torch.cuda.is_available() else "cpu"
hparams = TimesFmHparams(
backend=device,
per_core_batch_size=32,
horizon_len=128,
num_layers=50,
use_positional_embedding=False,
context_len=192,
)
if load_weights:
if FLAGS.local_model_path:
tfm_config = TimesFMConfig()
model = PatchedTimeSeriesDecoder(tfm_config)
loaded_checkpoint = load_file(FLAGS.local_model_path)
else:
repo_id = "google/timesfm-2.0-500m-pytorch"
tfm = TimesFm(hparams=hparams,
checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
tfm_config = tfm._model_config
model = PatchedTimeSeriesDecoder(tfm_config)
checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
model.load_state_dict(loaded_checkpoint)
return model, hparams, tfm_config
def plot_predictions(
model: TimesFm,
val_dataset: Dataset,
save_path: Optional[str] = "predictions.png",
) -> None:
"""
Plot model predictions against ground truth for a batch of validation data.
Args:
model: Trained TimesFM model
val_dataset: Validation dataset
save_path: Path to save the plot
"""
import matplotlib.pyplot as plt
model.eval()
x_context, x_padding, freq, x_future = val_dataset[0]
x_context = x_context.unsqueeze(0) # Add batch dimension
x_padding = x_padding.unsqueeze(0)
freq = freq.unsqueeze(0)
x_future = x_future.unsqueeze(0)
device = next(model.parameters()).device
x_context = x_context.to(device)
x_padding = x_padding.to(device)
freq = freq.to(device)
x_future = x_future.to(device)
with torch.no_grad():
predictions = model(x_context, x_padding.float(), freq)
predictions_mean = predictions[..., 0] # [B, N, horizon_len]
last_patch_pred = predictions_mean[:, -1, :] # [B, horizon_len]
context_vals = x_context[0].cpu().numpy()
future_vals = x_future[0].cpu().numpy()
pred_vals = last_patch_pred[0].cpu().numpy()
context_len = len(context_vals)
horizon_len = len(future_vals)
plt.figure(figsize=(12, 6))
plt.plot(range(context_len),
context_vals,
label="Historical Data",
color="blue",
linewidth=2)
plt.plot(
range(context_len, context_len + horizon_len),
future_vals,
label="Ground Truth",
color="green",
linestyle="--",
linewidth=2,
)
plt.plot(range(context_len, context_len + horizon_len),
pred_vals,
label="Prediction",
color="red",
linewidth=2)
plt.xlabel("Time Step")
plt.ylabel("Value")
plt.title("TimesFM Predictions vs Ground Truth")
plt.legend()
plt.grid(True)
if save_path:
plt.savefig(save_path)
print(f"Plot saved to {save_path}")
plt.close()
def get_data(context_len: int,
horizon_len: int,
freq_type: int = 0) -> Tuple[Dataset, Dataset]:
df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
time_series = df["Close"].values
train_dataset, val_dataset = prepare_datasets(
series=time_series,
context_length=context_len,
horizon_length=horizon_len,
freq_type=freq_type,
train_split=0.8,
)
print(f"Created datasets:")
print(f"- Training samples: {len(train_dataset)}")
print(f"- Validation samples: {len(val_dataset)}")
print(f"- Using frequency type: {freq_type}")
return train_dataset, val_dataset
def single_gpu_example():
"""Basic example of finetuning TimesFM on stock data."""
model, hparams, tfm_config = get_model(load_weights=True)
config = FinetuningConfig(batch_size=256,
num_epochs=5,
learning_rate=1e-4,
use_wandb=True,
freq_type=1,
log_every_n_steps=10,
val_check_interval=0.5,
use_quantile_loss=True)
train_dataset, val_dataset = get_data(128,
tfm_config.horizon_len,
freq_type=config.freq_type)
finetuner = TimesFMFinetuner(model, config)
print("\nStarting finetuning...")
results = finetuner.finetune(train_dataset=train_dataset,
val_dataset=val_dataset)
print("\nFinetuning completed!")
print(f"Training history: {len(results['history']['train_loss'])} epochs")
plot_predictions(
model=model,
val_dataset=val_dataset,
save_path="timesfm_predictions.png",
)
def setup_process(rank, world_size, model, config, train_dataset, val_dataset,
return_dict):
"""Setup process function with optimized CUDA handling."""
try:
if torch.cuda.is_available():
torch.cuda.set_device(rank)
os.environ["MASTER_ADDR"] = config.master_addr
os.environ["MASTER_PORT"] = config.master_port
if not torch.distributed.is_initialized():
torch.distributed.init_process_group(backend="nccl",
world_size=world_size,
rank=rank)
finetuner = TimesFMFinetuner(model, config, rank=rank)
results = finetuner.finetune(train_dataset=train_dataset,
val_dataset=val_dataset)
if rank == 0:
return_dict["results"] = results
plot_predictions(
model=model,
val_dataset=val_dataset,
save_path="timesfm_predictions.png",
)
except Exception as e:
print(f"Error in process {rank}: {str(e)}")
raise e
finally:
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
def multi_gpu_example():
"""Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
mp.set_start_method("spawn", force=True)
gpu_ids = [0, 1]
world_size = len(gpu_ids)
model, hparams, tfm_config = get_model(load_weights=True)
# Create config
config = FinetuningConfig(
batch_size=256,
num_epochs=5,
learning_rate=3e-5,
use_wandb=True,
distributed=True,
gpu_ids=gpu_ids,
log_every_n_steps=50,
val_check_interval=0.5,
)
train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
manager = mp.Manager()
return_dict = manager.dict()
# Launch processes
mp.spawn(
setup_process,
args=(world_size, model, config, train_dataset, val_dataset, return_dict),
nprocs=world_size,
join=True,
)
results = return_dict.get("results", None)
print("\nFinetuning completed!")
return results
def main(argv):
"""Main function that selects and runs the appropriate training mode."""
try:
if FLAGS.training_mode == "single":
print("\nStarting single-GPU training...")
single_gpu_example()
else:
gpu_ids = [int(id) for id in FLAGS.gpu_ids]
print(f"\nStarting multi-GPU training using GPUs: {gpu_ids}...")
config = FinetuningConfig(
batch_size=256,
num_epochs=5,
learning_rate=3e-5,
use_wandb=True,
distributed=True,
gpu_ids=gpu_ids,
)
results = multi_gpu_example(config)
print("\nMulti-GPU training completed!")
except Exception as e:
print(f"Training failed: {str(e)}")
finally:
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
if __name__ == "__main__":
app.run(main)
v1/src/finetuning/finetuning_torch.py
+399
View File
@@ -0,0 +1,399 @@
"""
TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
"""
import logging
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from timesfm.pytorch_patched_decoder import create_quantiles
import wandb
class MetricsLogger(ABC):
"""Abstract base class for logging metrics during training.
This class defines the interface for logging metrics during model training.
Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
"""
@abstractmethod
def log_metrics(self,
metrics: Dict[str, Any],
step: Optional[int] = None) -> None:
"""Log metrics to the specified backend.
Args:
metrics: Dictionary containing metric names and values.
step: Optional step number or epoch for the metrics.
"""
pass
@abstractmethod
def close(self) -> None:
"""Clean up any resources used by the logger."""
pass
class WandBLogger(MetricsLogger):
"""Weights & Biases implementation of metrics logging.
Args:
project: Name of the W&B project.
config: Configuration dictionary to log.
rank: Process rank in distributed training.
"""
def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
self.rank = rank
if rank == 0:
wandb.init(project=project, config=config)
def log_metrics(self,
metrics: Dict[str, Any],
step: Optional[int] = None) -> None:
"""Log metrics to W&B if on the main process.
Args:
metrics: Dictionary of metrics to log.
step: Current training step or epoch.
"""
if self.rank == 0:
wandb.log(metrics, step=step)
def close(self) -> None:
"""Finish the W&B run if on the main process."""
if self.rank == 0:
wandb.finish()
class DistributedManager:
"""Manages distributed training setup and cleanup.
Args:
world_size: Total number of processes.
rank: Process rank.
master_addr: Address of the master process.
master_port: Port for distributed communication.
backend: PyTorch distributed backend to use.
"""
def __init__(
self,
world_size: int,
rank: int,
master_addr: str = "localhost",
master_port: str = "12358",
backend: str = "nccl",
):
self.world_size = world_size
self.rank = rank
self.master_addr = master_addr
self.master_port = master_port
self.backend = backend
def setup(self) -> None:
"""Initialize the distributed environment."""
os.environ["MASTER_ADDR"] = self.master_addr
os.environ["MASTER_PORT"] = self.master_port
if not dist.is_initialized():
dist.init_process_group(backend=self.backend,
world_size=self.world_size,
rank=self.rank)
def cleanup(self) -> None:
"""Clean up the distributed environment."""
if dist.is_initialized():
dist.destroy_process_group()
@dataclass
class FinetuningConfig:
"""Configuration for model training.
Args:
batch_size: Number of samples per batch.
num_epochs: Number of training epochs.
learning_rate: Initial learning rate.
weight_decay: L2 regularization factor.
freq_type: Frequency, can be [0, 1, 2].
use_quantile_loss: bool = False # Flag to enable/disable quantile loss
quantiles: Optional[List[float]] = None
device: Device to train on ('cuda' or 'cpu').
distributed: Whether to use distributed training.
gpu_ids: List of GPU IDs to use.
master_port: Port for distributed training.
master_addr: Address for distributed training.
use_wandb: Whether to use Weights & Biases logging.
wandb_project: W&B project name.
log_every_n_steps: Log metrics every N steps (batches), this is inspired from Pytorch Lightning
val_check_interval: How often within one training epoch to check val metrics. (also from Pytorch Lightning)
Can be: float (0.0-1.0): fraction of epoch (e.g., 0.5 = validate twice per epoch)
int: validate every N batches
"""
batch_size: int = 32
num_epochs: int = 20
learning_rate: float = 1e-4
weight_decay: float = 0.01
freq_type: int = 0
use_quantile_loss: bool = False
quantiles: Optional[List[float]] = None
device: str = "cuda" if torch.cuda.is_available() else "cpu"
distributed: bool = False
gpu_ids: List[int] = field(default_factory=lambda: [0])
master_port: str = "12358"
master_addr: str = "localhost"
use_wandb: bool = False
wandb_project: str = "timesfm-finetuning"
log_every_n_steps: int = 50
val_check_interval: float = 0.5
class TimesFMFinetuner:
"""Handles model training and validation.
Args:
model: PyTorch model to train.
config: Training configuration.
rank: Process rank for distributed training.
loss_fn: Loss function (defaults to MSE).
logger: Optional logging.Logger instance.
"""
def __init__(
self,
model: nn.Module,
config: FinetuningConfig,
rank: int = 0,
loss_fn: Optional[Callable] = None,
logger: Optional[logging.Logger] = None,
):
self.model = model
self.config = config
self.rank = rank
self.logger = logger or logging.getLogger(__name__)
self.device = torch.device(
f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1))**2))
if config.use_wandb:
self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__,
rank)
if config.distributed:
self.dist_manager = DistributedManager(
world_size=len(config.gpu_ids),
rank=rank,
master_addr=config.master_addr,
master_port=config.master_port,
)
self.dist_manager.setup()
self.model = self._setup_distributed_model()
def _setup_distributed_model(self) -> nn.Module:
"""Configure model for distributed training."""
self.model = self.model.to(self.device)
return DDP(self.model,
device_ids=[self.config.gpu_ids[self.rank]],
output_device=self.config.gpu_ids[self.rank])
def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
"""Create appropriate DataLoader based on training configuration.
Args:
dataset: Dataset to create loader for.
is_train: Whether this is for training (affects shuffling).
Returns:
DataLoader instance.
"""
if self.config.distributed:
sampler = torch.utils.data.distributed.DistributedSampler(
dataset,
num_replicas=len(self.config.gpu_ids),
rank=dist.get_rank(),
shuffle=is_train)
else:
sampler = None
return DataLoader(
dataset,
batch_size=self.config.batch_size,
shuffle=(is_train and not self.config.distributed),
sampler=sampler,
)
def _quantile_loss(self, pred: torch.Tensor, actual: torch.Tensor,
quantile: float) -> torch.Tensor:
"""Calculates quantile loss.
Args:
pred: Predicted values
actual: Actual values
quantile: Quantile at which loss is computed
Returns:
Quantile loss
"""
dev = actual - pred
loss_first = dev * quantile
loss_second = -dev * (1.0 - quantile)
return 2 * torch.where(loss_first >= 0, loss_first, loss_second)
def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
"""Process a single batch of data.
Args:
batch: List of input tensors.
Returns:
Tuple of (loss, predictions).
"""
x_context, x_padding, freq, x_future = [
t.to(self.device, non_blocking=True) for t in batch
]
predictions = self.model(x_context, x_padding.float(), freq)
predictions_mean = predictions[..., 0]
last_patch_pred = predictions_mean[:, -1, :]
loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
if self.config.use_quantile_loss:
quantiles = self.config.quantiles or create_quantiles()
for i, quantile in enumerate(quantiles):
last_patch_quantile = predictions[:, -1, :, i + 1]
loss += torch.mean(
self._quantile_loss(last_patch_quantile, x_future.squeeze(-1),
quantile))
return loss, predictions
def _train_epoch(self, train_loader: DataLoader,
optimizer: torch.optim.Optimizer) -> float:
"""Train for one epoch in a distributed setting.
Args:
train_loader: DataLoader for training data.
optimizer: Optimizer instance.
Returns:
Average training loss for the epoch.
"""
self.model.train()
total_loss = 0.0
num_batches = len(train_loader)
for batch in train_loader:
loss, _ = self._process_batch(batch)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / num_batches
if self.config.distributed:
avg_loss_tensor = torch.tensor(avg_loss, device=self.device)
dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
avg_loss = (avg_loss_tensor / dist.get_world_size()).item()
return avg_loss
def _validate(self, val_loader: DataLoader) -> float:
"""Perform validation.
Args:
val_loader: DataLoader for validation data.
Returns:
Average validation loss.
"""
self.model.eval()
total_loss = 0.0
num_batches = len(val_loader)
with torch.no_grad():
for batch in val_loader:
loss, _ = self._process_batch(batch)
total_loss += loss.item()
avg_loss = total_loss / num_batches
if self.config.distributed:
avg_loss_tensor = torch.tensor(avg_loss, device=self.device)
dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
avg_loss = (avg_loss_tensor / dist.get_world_size()).item()
return avg_loss
def finetune(self, train_dataset: Dataset,
val_dataset: Dataset) -> Dict[str, Any]:
"""Train the model.
Args:
train_dataset: Training dataset.
val_dataset: Validation dataset.
Returns:
Dictionary containing training history.
"""
self.model = self.model.to(self.device)
train_loader = self._create_dataloader(train_dataset, is_train=True)
val_loader = self._create_dataloader(val_dataset, is_train=False)
optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.config.learning_rate,
weight_decay=self.config.weight_decay)
history = {"train_loss": [], "val_loss": [], "learning_rate": []}
self.logger.info(
f"Starting training for {self.config.num_epochs} epochs...")
self.logger.info(f"Training samples: {len(train_dataset)}")
self.logger.info(f"Validation samples: {len(val_dataset)}")
try:
for epoch in range(self.config.num_epochs):
train_loss = self._train_epoch(train_loader, optimizer)
val_loss = self._validate(val_loader)
current_lr = optimizer.param_groups[0]["lr"]
metrics = {
"train_loss": train_loss,
"val_loss": val_loss,
"learning_rate": current_lr,
"epoch": epoch + 1,
}
if self.config.use_wandb:
self.metrics_logger.log_metrics(metrics)
history["train_loss"].append(train_loss)
history["val_loss"].append(val_loss)
history["learning_rate"].append(current_lr)
if self.rank == 0:
self.logger.info(
f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}"
)
except KeyboardInterrupt:
self.logger.info("Training interrupted by user")
if self.config.distributed:
self.dist_manager.cleanup()
if self.config.use_wandb:
self.metrics_logger.close()
return {"history": history}
v1/src/timesfm/__init__.py
+35
View File
@@ -0,0 +1,35 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TimesFM init file."""
print(
" See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs."
)
from timesfm.timesfm_base import (
freq_map,
TimesFmCheckpoint,
TimesFmHparams,
TimesFmBase,
)
import sys
try:
from timesfm.timesfm_jax import TimesFmJax as TimesFm
from timesfm import data_loader
print(f"Loaded Jax TimesFM, likely because python version is {sys.version}.")
except Exception as _:
from timesfm.timesfm_torch import TimesFmTorch as TimesFm
print(f"Loaded PyTorch TimesFM, likely because python version is {sys.version}.")
v1/src/timesfm/data_loader.py
+255
View File
@@ -0,0 +1,255 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TF dataloaders for general timeseries datasets.
The expected input format is csv file with a datetime index.
"""
from absl import logging
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from . import time_features
class TimeSeriesdata(object):
"""Data loader class."""
def __init__(
self,
data_path,
datetime_col,
num_cov_cols,
cat_cov_cols,
ts_cols,
train_range,
val_range,
test_range,
hist_len,
pred_len,
batch_size,
freq='H',
normalize=True,
epoch_len=None,
holiday=False,
permute=True,
):
"""Initialize objects.
Args:
data_path: path to csv file
datetime_col: column name for datetime col
num_cov_cols: list of numerical global covariates
cat_cov_cols: list of categorical global covariates
ts_cols: columns corresponding to ts
train_range: tuple of train ranges
val_range: tuple of validation ranges
test_range: tuple of test ranges
hist_len: historical context
pred_len: prediction length
batch_size: batch size (number of ts in a batch)
freq: freq of original data
normalize: std. normalize data or not
epoch_len: num iters in an epoch
holiday: use holiday features or not
permute: permute ts in train batches or not
Returns:
None
"""
self.data_df = pd.read_csv(open(data_path, 'r'))
if not num_cov_cols:
self.data_df['ncol'] = np.zeros(self.data_df.shape[0])
num_cov_cols = ['ncol']
if not cat_cov_cols:
self.data_df['ccol'] = np.zeros(self.data_df.shape[0])
cat_cov_cols = ['ccol']
self.data_df.fillna(0, inplace=True)
self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]),
inplace=True)
self.num_cov_cols = num_cov_cols
self.cat_cov_cols = cat_cov_cols
self.ts_cols = ts_cols
self.train_range = train_range
self.val_range = val_range
self.test_range = test_range
data_df_idx = self.data_df.index
date_index = data_df_idx.union(
pd.date_range(
data_df_idx[-1] + pd.Timedelta(1, freq=freq),
periods=pred_len + 1,
freq=freq,
))
self.time_df = time_features.TimeCovariates(
date_index, holiday=holiday).get_covariates()
self.hist_len = hist_len
self.pred_len = pred_len
self.batch_size = batch_size
self.freq = freq
self.normalize = normalize
self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose()
self.data_mat = self.data_mat[:, 0:self.test_range[1]]
self.time_mat = self.time_df.to_numpy().transpose()
self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose()
self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols)
self.normalize = normalize
if normalize:
self._normalize_data()
logging.info(
'Data Shapes: %s, %s, %s, %s',
self.data_mat.shape,
self.time_mat.shape,
self.num_feat_mat.shape,
self.cat_feat_mat.shape,
)
self.epoch_len = epoch_len
self.permute = permute
def _get_cat_cols(self, cat_cov_cols):
"""Get categorical columns."""
cat_vars = []
cat_sizes = []
for col in cat_cov_cols:
dct = {x: i for i, x in enumerate(self.data_df[col].unique())}
cat_sizes.append(len(dct))
mapped = self.data_df[col].map(lambda x: dct[x]).to_numpy().transpose() # pylint: disable=cell-var-from-loop
cat_vars.append(mapped)
return np.vstack(cat_vars), cat_sizes
def _normalize_data(self):
self.scaler = StandardScaler()
train_mat = self.data_mat[:, 0:self.train_range[1]]
self.scaler = self.scaler.fit(train_mat.transpose())
self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose()
def train_gen(self):
"""Generator for training data."""
num_ts = len(self.ts_cols)
perm = np.arange(
self.train_range[0] + self.hist_len,
self.train_range[1] - self.pred_len,
)
perm = np.random.permutation(perm)
hist_len = self.hist_len
logging.info('Hist len: %s', hist_len)
if not self.epoch_len:
epoch_len = len(perm)
else:
epoch_len = self.epoch_len
for idx in perm[0:epoch_len]:
for _ in range(num_ts // self.batch_size + 1):
if self.permute:
tsidx = np.random.choice(num_ts, size=self.batch_size, replace=False)
else:
tsidx = np.arange(num_ts)
dtimes = np.arange(idx - hist_len, idx + self.pred_len)
(
bts_train,
bts_pred,
bfeats_train,
bfeats_pred,
bcf_train,
bcf_pred,
) = self._get_features_and_ts(dtimes, tsidx, hist_len)
all_data = [
bts_train,
bfeats_train,
bcf_train,
bts_pred,
bfeats_pred,
bcf_pred,
tsidx,
]
yield tuple(all_data)
def test_val_gen(self, mode='val', shift=1):
"""Generator for validation/test data."""
if mode == 'val':
start = self.val_range[0]
end = self.val_range[1] - self.pred_len + 1
elif mode == 'test':
start = self.test_range[0]
end = self.test_range[1] - self.pred_len + 1
else:
raise NotImplementedError('Eval mode not implemented')
num_ts = len(self.ts_cols)
hist_len = self.hist_len
logging.info('Hist len: %s', hist_len)
perm = np.arange(start, end)
if self.epoch_len:
epoch_len = self.epoch_len
else:
epoch_len = len(perm)
for i in range(0, epoch_len, shift):
idx = perm[i]
for batch_idx in range(0, num_ts, self.batch_size):
tsidx = np.arange(batch_idx, min(batch_idx + self.batch_size, num_ts))
dtimes = np.arange(idx - hist_len, idx + self.pred_len)
(
bts_train,
bts_pred,
bfeats_train,
bfeats_pred,
bcf_train,
bcf_pred,
) = self._get_features_and_ts(dtimes, tsidx, hist_len)
all_data = [
bts_train,
bfeats_train,
bcf_train,
bts_pred,
bfeats_pred,
bcf_pred,
tsidx,
]
yield tuple(all_data)
def _get_features_and_ts(self, dtimes, tsidx, hist_len=None):
"""Get features and ts in specified windows."""
if hist_len is None:
hist_len = self.hist_len
data_times = dtimes[dtimes < self.data_mat.shape[1]]
bdata = self.data_mat[:, data_times]
bts = bdata[tsidx, :]
bnf = self.num_feat_mat[:, data_times]
bcf = self.cat_feat_mat[:, data_times]
btf = self.time_mat[:, dtimes]
if bnf.shape[1] < btf.shape[1]:
rem_len = btf.shape[1] - bnf.shape[1]
rem_rep = np.repeat(bnf[:, [-1]], repeats=rem_len)
rem_rep_cat = np.repeat(bcf[:, [-1]], repeats=rem_len)
bnf = np.hstack([bnf, rem_rep.reshape(bnf.shape[0], -1)])
bcf = np.hstack([bcf, rem_rep_cat.reshape(bcf.shape[0], -1)])
bfeats = np.vstack([btf, bnf])
bts_train = bts[:, 0:hist_len]
bts_pred = bts[:, hist_len:]
bfeats_train = bfeats[:, 0:hist_len]
bfeats_pred = bfeats[:, hist_len:]
bcf_train = bcf[:, 0:hist_len]
bcf_pred = bcf[:, hist_len:]
return bts_train, bts_pred, bfeats_train, bfeats_pred, bcf_train, bcf_pred
def tf_dataset(self, mode='train', shift=1):
"""Tensorflow Dataset."""
if mode == 'train':
gen_fn = self.train_gen
else:
gen_fn = lambda: self.test_val_gen(mode, shift)
output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 +
[tf.int32] * 2)
dataset = tf.data.Dataset.from_generator(gen_fn, output_types)
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
return dataset
v1/src/timesfm/patched_decoder.py
+543
View File
@@ -0,0 +1,543 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Pax ML model for patched time-series decoder.
The file implements Residual MLPs, Patched Decoder layers and PAX ML models.
"""
import dataclasses
from typing import Optional, Tuple
import einshape as es
from jax import lax
import jax.numpy as jnp
from praxis import base_layer
from praxis import base_model
from praxis import layers
from praxis import pax_fiddle
from praxis import py_utils
from praxis import pytypes
from praxis.layers import activations
from praxis.layers import embedding_softmax
from praxis.layers import linears
from praxis.layers import normalizations
from praxis.layers import stochastics
from praxis.layers import transformers
# PAX shortcuts
NestedMap = py_utils.NestedMap
JTensor = pytypes.JTensor
LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]
template_field = base_layer.template_field
PAD_VAL = 1123581321.0
DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
# NestedMap keys
_INPUT_TS = "input_ts"
_TARGET_FUTURE = "actual_ts"
_INPUT_PADDING = "input_padding"
_OUTPUT_TS = "output_ts"
_FREQ = "freq"
_OUTPUT_TOKENS = "output_tokens"
_STATS = "stats"
# Small numerical value.
_TOLERANCE = 1e-7
def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor:
"""Shifts rows of seq based on the first 0 in each row of the mask."""
num = seq.shape[1]
# Find the index of the first 0 in each row of the mask
first_zero_idx = jnp.argmin(mask, axis=1)
# Create a range array for indexing
idx_range = jnp.arange(num)
def shift_row(carry, x):
seq_row, shift = x
shifted_idx = (idx_range - shift) % num
shifted_row = seq_row[shifted_idx]
return carry, shifted_row
# Use lax.scan to shift each row of seq based on the corresponding
# first_zero_idx.
_, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx))
return shifted_seq
class ResidualBlock(base_layer.BaseLayer):
"""Simple feedforward block with residual connection.
Attributes:
input_dims: input dimension.
hidden_dims: hidden dimension.
output_dims: output dimension.
dropout_prob: dropout probability.
layer_norm: whether to use layer norm or not.
dropout_tpl: config for dropout.
ln_tpl: config for layer norm.
act_tpl: config for activation in hidden layer.
"""
input_dims: int = 0
hidden_dims: int = 0
output_dims: int = 0
dropout_prob: float = 0.0
layer_norm: bool = False
dropout_tpl: LayerTpl = template_field(stochastics.Dropout)
ln_tpl: LayerTpl = template_field(normalizations.LayerNorm)
act_tpl: LayerTpl = template_field(activations.Swish)
def setup(self):
lnorm_tpl = self.ln_tpl.clone()
lnorm_tpl.dim = self.output_dims
self.create_child("ln_layer", lnorm_tpl)
dropout_tpl = self.dropout_tpl.clone()
dropout_tpl.keep_prob = 1.0 - self.dropout_prob
self.create_child("dropout", dropout_tpl)
self.create_child(
"hidden_layer",
pax_fiddle.Config(
linears.FeedForward,
input_dims=self.input_dims,
output_dims=self.hidden_dims,
activation_tpl=self.act_tpl.clone(),
),
)
self.create_child(
"output_layer",
pax_fiddle.Config(
linears.FeedForward,
input_dims=self.hidden_dims,
output_dims=self.output_dims,
activation_tpl=pax_fiddle.Config(activations.Identity),
),
)
self.create_child(
"residual_layer",
pax_fiddle.Config(
linears.FeedForward,
input_dims=self.input_dims,
output_dims=self.output_dims,
activation_tpl=pax_fiddle.Config(activations.Identity),
),
)
def __call__(self, inputs: JTensor) -> JTensor:
hidden = self.hidden_layer(inputs)
output = self.output_layer(hidden)
output = self.dropout(output)
residual = self.residual_layer(inputs)
if self.layer_norm:
return self.ln_layer(output + residual)
else:
return output + residual
def _masked_mean_std(inputs: JTensor,
padding: JTensor) -> Tuple[JTensor, JTensor]:
"""Calculates mean and standard deviation of arr across axis 1.
It should exclude values where pad is 1.
Args:
inputs: A JAX array of shape [b, n, p].
padding: A JAX array of shape [b, n, p] with values 0 or 1.
Returns:
A tuple containing the mean and standard deviation of arr. We return the
statistics of the first patch with more than three non-padded values.
"""
# Selecting the first pad with more than 3 unpadded values.
pad_sum = jnp.sum(1 - padding, axis=2)
def _get_patch_index(arr: JTensor):
indices = jnp.argmax(arr >= 3, axis=1)
row_sum = (arr >= 3).sum(axis=1)
return jnp.where(row_sum == 0, arr.shape[1] - 1, indices)
patch_indices = _get_patch_index(pad_sum)
bidxs = jnp.arange(inputs.shape[0])
arr = inputs[bidxs, patch_indices, :]
pad = padding[bidxs, patch_indices, :]
# Create a mask where P is 0
mask = 1 - pad
# Calculate the number of valid elements
num_valid_elements = jnp.sum(mask, axis=1)
num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements)
# Calculate the masked sum and squared sum of M
masked_sum = jnp.sum(arr * mask, axis=1)
masked_squared_sum = jnp.sum((arr * mask)**2, axis=1)
# Calculate the masked mean and standard deviation
masked_mean = masked_sum / num_valid_elements
masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var)
masked_std = jnp.sqrt(masked_var)
return masked_mean, masked_std
def _create_quantiles() -> list[float]:
"""Returns the quantiles for forecasting."""
return DEFAULT_QUANTILES
class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
"""Patch decoder layer for time-series foundation model.
Attributes:
patch_len: length of input patches.
horizon_len: length of output patches. Referred to as `output_patch_len`
during inference.
model_dims: model dimension of stacked transformer layer.
hidden_dims: hidden dimensions in fully connected layers.
quantiles: list of quantiles for non prob model.
residual_block_tpl: config for residual block.
stacked_transformer_params_tpl: config for stacked transformer.
use_freq: whether to use frequency encoding.
In all of what followed, except specified otherwise, B is batch size, T is
sequence length of time-series. N is the number of input patches that can be
obtained from T. P is the input patch length and H is the horizon length. Q is
number of output logits. D is model dimension.
"""
patch_len: int = 0
horizon_len: int = 0
model_dims: int = 0
hidden_dims: int = 0
quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles)
residual_block_tpl: LayerTpl = template_field(ResidualBlock)
stacked_transformer_params_tpl: LayerTpl = template_field(
transformers.StackedTransformer)
use_freq: bool = True
use_pos_emb: bool = True
def setup(self) -> None:
"""Construct the model."""
num_outputs = len(self.quantiles) + 1
stl = self.stacked_transformer_params_tpl.clone()
stl.model_dims = self.model_dims
stl.hidden_dims = self.hidden_dims
stl.mask_self_attention = True
self.create_child("stacked_transformer_layer", stl)
input_resl = self.residual_block_tpl.clone()
ff_in_dims = 2 * self.patch_len
input_resl.input_dims = ff_in_dims
input_resl.hidden_dims = self.hidden_dims
input_resl.output_dims = self.model_dims
self.create_child(
"input_ff_layer",
input_resl,
)
horizon_resl = self.residual_block_tpl.clone()
horizon_resl.input_dims = self.model_dims
horizon_resl.hidden_dims = self.hidden_dims
horizon_resl.output_dims = self.horizon_len * num_outputs
self.create_child(
"horizon_ff_layer",
horizon_resl,
)
self.create_child(
"position_emb",
pax_fiddle.Config(layers.PositionalEmbedding,
embedding_dims=self.model_dims),
)
if self.use_freq:
self.create_child(
"freq_emb",
pax_fiddle.Config(
embedding_softmax.Embedding,
num_classes=3,
input_dims=self.model_dims,
),
)
def transform_decode_state(
self, transform_fn: base_layer.DecodeStateTransformFn) -> None:
"""Transforms all decode state variables based on transform_fn."""
self.stacked_transformer_layer.transform_decode_state(transform_fn)
def _forward_transform(
self, inputs: JTensor,
patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
"""Input is of shape [B, N, P]."""
mu, sigma = _masked_mean_std(inputs, patched_pads)
sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma)
# Normalize each patch.
outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
outputs = jnp.where(
jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs)
return outputs, (mu, sigma)
def _reverse_transform(self, outputs: JTensor,
stats: Tuple[JTensor, JTensor]) -> JTensor:
"""Output is of shape [B, N, P, Q]."""
mu, sigma = stats
return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
def _preprocess_input(
self,
input_ts: JTensor,
input_padding: JTensor,
pos_emb: Optional[JTensor] = None,
) -> Tuple[JTensor, JTensor, Optional[Tuple[JTensor, JTensor]], JTensor]:
"""Preprocess input for stacked transformer."""
# Reshape into patches.
patched_inputs = es.jax_einshape("b(np)->bnp", input_ts, p=self.patch_len)
patched_pads = es.jax_einshape("b(np)->bnp",
input_padding,
p=self.patch_len)
patched_inputs = jnp.where(
jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs)
patched_pads = jnp.where(
jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads)
patched_inputs, stats = self._forward_transform(patched_inputs,
patched_pads)
# B x N x D
patched_inputs = patched_inputs * (1.0 - patched_pads)
concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1)
model_input = self.input_ff_layer(concat_inputs)
# A patch should not be padded even if there is at least one zero.
patched_padding = jnp.min(patched_pads, axis=-1)
if self.use_pos_emb:
if pos_emb is None:
position_emb = self.position_emb(seq_length=model_input.shape[1])
else:
position_emb = pos_emb
if self.do_eval:
if position_emb.shape[0] != model_input.shape[0]:
position_emb = jnp.repeat(position_emb, model_input.shape[0], axis=0)
position_emb = _shift_padded_seq(patched_padding, position_emb)
model_input += position_emb
return model_input, patched_padding, stats, patched_inputs
def _postprocess_output(
self,
model_output: JTensor,
num_outputs: int,
stats: Tuple[JTensor, JTensor],
) -> JTensor:
"""Postprocess output of stacked transformer."""
# B x N x (H.Q)
output_ts = self.horizon_ff_layer(model_output)
output_ts = es.jax_einshape("bn(hq)->bnhq",
output_ts,
q=num_outputs,
h=self.horizon_len)
return self._reverse_transform(output_ts, stats)
def __call__(self, inputs: NestedMap) -> NestedMap:
"""PatchTST call.
Args:
inputs: A NestedMap containing (1) input_ts: input sequence of shape [B,
T] where T must be multiple of patch_length; (2) input_padding: that
contains padding map.
Returns:
A nested map with two keys:
(1) 'output_tokens' of shape [B, N, D].
(2) 'output_ts' of shape [B, N, H, Q]
(3) 'stats' a Tuple of statistics for renormalization.
"""
input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING]
num_outputs = len(self.quantiles) + 1
model_input, patched_padding, stats, _ = self._preprocess_input(
input_ts=input_ts,
input_padding=input_padding,
)
if self.use_freq:
freq = inputs[_FREQ].astype(jnp.int32)
f_emb = self.freq_emb(freq) # B x 1 x D
f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1)
model_input += f_emb
model_output = self.stacked_transformer_layer(model_input, patched_padding)
output_ts = self._postprocess_output(model_output, num_outputs, stats)
return NestedMap({
_OUTPUT_TOKENS: model_output,
_OUTPUT_TS: output_ts,
_STATS: stats
})
def decode(
self,
inputs: NestedMap,
horizon_len: int,
output_patch_len: Optional[int] = None,
max_len: int | None = None,
return_forecast_on_context: bool = False,
) -> tuple[JTensor, JTensor]:
"""Auto-regressive decoding without caching.
Args:
inputs: input time-series and paddings. Time-series shape B x C, padding
shape shape B x (C + H) where H is the prediction length.
horizon_len: prediction length.
output_patch_len: output length to be fetched from one step of
auto-regressive decoding.
max_len: maximum training context length.
return_forecast_on_context: whether to return the model forecast on the
context except the first input patch.
Returns:
Tuple of two forecasting results:
- Point (mean) output predictions as a tensor with shape B x H'.
- Full predictions (mean and quantiles) as a tensor with shape
B x H' x (1 + # quantiles).
In particular, if return_forecast_on_context is True, H' is H plus
the forecastable context length, i.e. context_len - (first) patch_len.
"""
final_out = inputs[_INPUT_TS]
context_len = final_out.shape[1]
paddings = inputs[_INPUT_PADDING]
if max_len is None:
max_len = context_len
if self.use_freq:
freq = inputs[_FREQ].astype(jnp.int32)
else:
freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32)
full_outputs = []
if paddings.shape[1] != final_out.shape[1] + horizon_len:
raise ValueError(
"Length of paddings must match length of input + horizon_len:"
f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}")
if output_patch_len is None:
output_patch_len = self.horizon_len
num_decode_patches = (horizon_len + output_patch_len -
1) // output_patch_len
for step_index in range(num_decode_patches):
current_padding = paddings[:, 0:final_out.shape[1]]
input_ts = final_out[:, -max_len:]
input_padding = current_padding[:, -max_len:]
model_input = NestedMap(
input_ts=input_ts,
input_padding=input_padding,
freq=freq,
)
fprop_outputs = self(model_input)[_OUTPUT_TS]
if return_forecast_on_context and step_index == 0:
# For the first decodings step, collect the model forecast on the
# context except the unavailable first input batch forecast.
new_full_ts = fprop_outputs[:, :-1, :self.patch_len, :]
new_full_ts = es.jax_einshape("bnph->b(np)h", new_full_ts)
full_outputs.append(new_full_ts)
# (full batch, last patch, output_patch_len, index of mean forecast = 0)
new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
new_full_ts = fprop_outputs[:, -1, :output_patch_len, :]
# (full batch, last patch, output_patch_len, all output indices)
full_outputs.append(new_full_ts)
final_out = jnp.concatenate([final_out, new_ts], axis=-1)
if return_forecast_on_context:
# `full_outputs` indexing starts at after the first input patch.
full_outputs = jnp.concatenate(full_outputs,
axis=1)[:, :(context_len - self.patch_len +
horizon_len), :]
else:
# `full_outputs` indexing starts at the forecast horizon.
full_outputs = jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :]
return (full_outputs[:, :, 0], full_outputs)
class PatchedDecoderFinetuneModel(base_model.BaseModel):
"""Model class for finetuning patched time-series decoder.
Attributes:
core_layer_tpl: config for core layer.
freq: freq to finetune on.
"""
core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder)
freq: int = 0
def setup(self) -> None:
self.create_child("core_layer", self.core_layer_tpl)
def compute_predictions(self, input_batch: NestedMap) -> NestedMap:
input_ts = input_batch[_INPUT_TS]
input_padding = jnp.zeros_like(input_ts)
context_len = input_ts.shape[1]
input_patch_len = self.core_layer_tpl.patch_len
context_pad = ((context_len + input_patch_len - 1) //
input_patch_len) * input_patch_len - context_len
input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)])
input_padding = jnp.pad(input_padding, [(0, 0), (context_pad, 0)],
constant_values=1)
freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
new_input_batch = NestedMap(
input_ts=input_ts,
input_padding=input_padding,
freq=freq,
)
return self.core_layer(new_input_batch)
def _quantile_loss(self, pred: JTensor, actual: JTensor,
quantile: float) -> JTensor:
"""Calculates quantile loss.
Args:
pred: B x T
actual: B x T
quantile: quantile at which loss is computed.
Returns:
per coordinate loss.
"""
dev = actual - pred
loss_first = dev * quantile
loss_second = -dev * (1.0 - quantile)
return 2 * jnp.where(loss_first >= 0, loss_first, loss_second)
def compute_loss(self, prediction_output: NestedMap,
input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]:
output_ts = prediction_output[_OUTPUT_TS]
actual_ts = input_batch[_TARGET_FUTURE]
pred_ts = output_ts[:, -1, 0:actual_ts.shape[1], :]
loss = jnp.square(pred_ts[:, :, 0] - actual_ts)
for i, quantile in enumerate(self.core_layer.quantiles):
loss += self._quantile_loss(pred_ts[:, :, i + 1], actual_ts, quantile)
loss = loss.mean()
loss_weight = jnp.array(1.0, dtype=jnp.float32)
per_example_out = NestedMap()
return {"avg_qloss": (loss, loss_weight)}, per_example_out
v1/src/timesfm/pytorch_patched_decoder.py
+801
View File
@@ -0,0 +1,801 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Pytorch version of patched decoder."""
import dataclasses
import math
from typing import List, Tuple
import torch
from torch import nn
import torch.nn.functional as F
def create_quantiles() -> list[float]:
return [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
@dataclasses.dataclass
class TimesFMConfig:
"""Config for initializing timesfm patched_decoder class."""
# The number of blocks in the model.
num_layers: int = 20
# The number of attention heads used in the attention layers of the model.
num_heads: int = 16
# The number of key-value heads for implementing attention.
num_kv_heads: int = 16
# The hidden size of the model.
hidden_size: int = 1280
# The dimension of the MLP representations.
intermediate_size: int = 1280
# The number of head dimensions.
head_dim: int = 80
# The epsilon used by the rms normalization layers.
rms_norm_eps: float = 1e-6
# Patch length
patch_len: int = 32
# Horizon length
horizon_len: int = 128
# quantiles
quantiles: List[float] = dataclasses.field(default_factory=create_quantiles)
# Padding value
pad_val: float = 1123581321.0
# Tolerance
tolerance: float = 1e-6
# The dtype of the weights.
dtype: str = "bfloat32"
# use positional embedding
use_positional_embedding: bool = True
def _masked_mean_std(
inputs: torch.Tensor,
padding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Calculates mean and standard deviation of `inputs` across axis 1.
It excludes values where `padding` is 1.
Args:
inputs: A PyTorch tensor of shape [b, n, p].
padding: A PyTorch tensor of shape [b, n, p] with values 0 or 1.
Returns:
A tuple containing the mean and standard deviation.
We return the statistics of the first patch with more than three non-padded
values.
"""
# Selecting the first patch with more than 3 unpadded values.
pad_sum = torch.sum(1 - padding, dim=2)
def _get_patch_index(arr: torch.Tensor):
indices = torch.argmax((arr >= 3).to(torch.int32), dim=1)
row_sum = (arr >= 3).to(torch.int32).sum(dim=1)
return torch.where(row_sum == 0, arr.shape[1] - 1, indices)
patch_indices = _get_patch_index(pad_sum)
bidxs = torch.arange(inputs.shape[0])
arr = inputs[bidxs, patch_indices, :]
pad = padding[bidxs, patch_indices, :]
# Create a mask where padding is 0
mask = 1 - pad
# Calculate the number of valid elements
num_valid_elements = torch.sum(mask, dim=1)
num_valid_elements = torch.where(
num_valid_elements == 0,
torch.tensor(1,
dtype=num_valid_elements.dtype,
device=num_valid_elements.device),
num_valid_elements,
)
# Calculate the masked sum and squared sum
masked_sum = torch.sum(arr * mask, dim=1)
masked_squared_sum = torch.sum((arr * mask)**2, dim=1)
# Calculate the masked mean and standard deviation
masked_mean = masked_sum / num_valid_elements
masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
masked_var = torch.where(
masked_var < 0.0,
torch.tensor(0.0, dtype=masked_var.dtype, device=masked_var.device),
masked_var,
)
masked_std = torch.sqrt(masked_var)
return masked_mean, masked_std
def _shift_padded_seq(mask: torch.Tensor, seq: torch.Tensor) -> torch.Tensor:
"""Shifts rows of seq based on the first 0 in each row of the mask.
Args:
mask: mask tensor of shape [B, N]
seq: seq tensor of shape [B, N, P]
Returns:
Returns the shifted sequence.
"""
batch_size, num_seq, feature_dim = seq.shape
new_mask: torch.BoolTensor = mask == 0
# Use argmax to find the first True value in each row
indices = new_mask.to(torch.int32).argmax(dim=1)
# Handle rows with all zeros
indices[~new_mask.any(dim=1)] = -1
# Create index ranges for each sequence in the batch
idx_range = (torch.arange(num_seq).to(
seq.device).unsqueeze(0).unsqueeze(-1).expand(batch_size, -1,
feature_dim))
# Calculate shifted indices for each element in each sequence
shifted_idx = (idx_range - indices[:, None, None]) % num_seq
# Gather values from seq using shifted indices
shifted_seq = seq.gather(1, shifted_idx)
return shifted_seq
def get_large_negative_number(dtype: torch.dtype) -> torch.Tensor:
"""Returns a large negative value for the given dtype."""
if dtype.is_floating_point:
dtype_max = torch.finfo(dtype).max
else:
dtype_max = torch.iinfo(dtype).max
return torch.tensor(-0.7 * dtype_max, dtype=dtype)
def apply_mask_to_logits(logits: torch.Tensor,
mask: torch.Tensor) -> torch.Tensor:
"""Applies a floating-point mask to a set of logits.
Args:
logits: A torch.Tensor of logit values.
mask: A torch.Tensor (float32) of mask values with the encoding described
in the function documentation.
Returns:
Masked logits.
"""
min_value = get_large_negative_number(logits.dtype)
return torch.where((mask >= min_value * 0.5), logits, min_value)
def convert_paddings_to_mask(
paddings: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
"""Converts binary paddings to a logit mask ready to add to attention matrix.
Args:
paddings: binary torch.Tensor of shape [B, T], with 1 denoting padding
token.
dtype: data type of the input.
Returns:
A torch.Tensor of shape [B, 1, 1, T] ready to add to attention logits.
"""
attention_mask = paddings.detach().clone()
attention_mask = attention_mask[:, None, None, :] # Equivalent to jnp.newaxis
attention_mask *= get_large_negative_number(dtype)
return attention_mask
def causal_mask(input_t: torch.Tensor) -> torch.Tensor:
"""Computes and returns causal mask.
Args:
input_t: A torch.Tensor of shape [B, T, D].
Returns:
An attention_mask torch.Tensor of shape [1, 1, T, T]. Attention mask has
already been converted to large negative values.
"""
assert input_t.dtype.is_floating_point, input_t.dtype
large_negative_number = get_large_negative_number(input_t.dtype)
t = input_t.shape[1]
col_idx = torch.arange(t).unsqueeze(0).repeat(t, 1)
row_idx = torch.arange(t).unsqueeze(1).repeat(1, t)
mask = (row_idx < col_idx).to(input_t.dtype) * large_negative_number
return (mask.unsqueeze(0).unsqueeze(0).to(input_t.device)
) # Equivalent to jnp.newaxis
def merge_masks(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""Merges 2 masks.
logscale mask is expected but 0/1 mask is also fine.
Args:
a: torch.Tensor of shape [1|B, 1, 1|T, S].
b: torch.Tensor of shape [1|B, 1, 1|T, S].
Returns:
torch.Tensor of shape [1|B, 1, 1|T, S].
"""
def expand_t(key_mask):
query_mask = key_mask.transpose(-1, -2) # Equivalent of jnp.transpose
return torch.minimum(query_mask, key_mask)
if a.shape[2] != b.shape[2]:
if a.shape[2] == 1:
a = expand_t(a)
else:
assert b.shape[2] == 1
b = expand_t(b)
assert a.shape[1:] == b.shape[1:], f"a.shape={a.shape}, b.shape={b.shape}."
return torch.minimum(a, b) # Element-wise minimum, similar to jnp.minimum
class ResidualBlock(nn.Module):
"""TimesFM residual block."""
def __init__(
self,
input_dims,
hidden_dims,
output_dims,
):
super(ResidualBlock, self).__init__()
self.input_dims = input_dims
self.hidden_dims = hidden_dims
self.output_dims = output_dims
# Hidden Layer
self.hidden_layer = nn.Sequential(
nn.Linear(input_dims, hidden_dims),
nn.SiLU(),
)
# Output Layer
self.output_layer = nn.Linear(hidden_dims, output_dims)
# Residual Layer
self.residual_layer = nn.Linear(input_dims, output_dims)
def forward(self, x):
hidden = self.hidden_layer(x)
output = self.output_layer(hidden)
residual = self.residual_layer(x)
return output + residual
class RMSNorm(torch.nn.Module):
"""Pax rms norm in pytorch."""
def __init__(
self,
dim: int,
eps: float = 1e-6,
add_unit_offset: bool = False,
):
super().__init__()
self.eps = eps
self.add_unit_offset = add_unit_offset
self.weight = nn.Parameter(torch.zeros(dim))
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
output = self._norm(x.float())
if self.add_unit_offset:
output = output * (1 + self.weight.float())
else:
output = output * self.weight.float()
return output.type_as(x)
class TransformerMLP(nn.Module):
"""Pax transformer MLP in pytorch."""
def __init__(
self,
hidden_size: int,
intermediate_size: int,
):
super().__init__()
self.gate_proj = nn.Linear(hidden_size, intermediate_size)
self.down_proj = nn.Linear(intermediate_size, hidden_size)
self.layer_norm = nn.LayerNorm(normalized_shape=hidden_size, eps=1e-6)
def forward(self, x, paddings=None):
gate_inp = self.layer_norm(x)
gate = self.gate_proj(gate_inp)
gate = F.relu(gate)
outputs = self.down_proj(gate)
if paddings is not None:
outputs = outputs * (1.0 - paddings[:, :, None])
return outputs + x
class TimesFMAttention(nn.Module):
"""Implements the attention used in TimesFM."""
def __init__(
self,
hidden_size: int,
num_heads: int,
num_kv_heads: int,
head_dim: int,
):
super().__init__()
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
assert self.num_heads % self.num_kv_heads == 0
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
self.hidden_size = hidden_size
self.head_dim = head_dim
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.scaling = nn.Parameter(
torch.empty((self.head_dim,), dtype=torch.float32),)
self.qkv_proj = nn.Linear(
self.hidden_size,
(self.num_heads + 2 * self.num_kv_heads) * self.head_dim,
)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size)
def _per_dim_scaling(self, query: torch.Tensor) -> torch.Tensor:
# [batch_size, n_local_heads, input_len, head_dim]
r_softplus_0 = 1.442695041
softplus_func = torch.nn.Softplus()
scale = r_softplus_0 / math.sqrt(self.head_dim)
scale = scale * softplus_func(self.scaling)
return query * scale[None, None, None, :]
def forward(
self,
hidden_states: torch.Tensor,
mask: torch.Tensor,
kv_write_indices: torch.Tensor | None = None,
kv_cache: Tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
hidden_states_shape = hidden_states.shape
assert len(hidden_states_shape) == 3
batch_size, input_len, _ = hidden_states_shape
qkv = self.qkv_proj(hidden_states)
xq, xk, xv = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
xq = xq.view(batch_size, -1, self.num_heads, self.head_dim)
xk = xk.view(batch_size, -1, self.num_kv_heads, self.head_dim)
xv = xv.view(batch_size, -1, self.num_kv_heads, self.head_dim)
xq = self._per_dim_scaling(xq)
# Write new kv cache.
# [batch_size, input_len, n_local_kv_heads, head_dim]
if kv_cache is not None and kv_write_indices is not None:
k_cache, v_cache = kv_cache
k_cache.index_copy_(1, kv_write_indices, xk)
v_cache.index_copy_(1, kv_write_indices, xv)
key = k_cache
value = v_cache
else:
key = xk
value = xv
if self.num_kv_heads != self.num_heads:
# [batch_size, max_seq_len, n_local_heads, head_dim]
key = torch.repeat_interleave(key, self.num_queries_per_kv, dim=2)
value = torch.repeat_interleave(value, self.num_queries_per_kv, dim=2)
# [batch_size, n_local_heads, input_len, head_dim]
q = xq.transpose(1, 2)
# [batch_size, n_local_heads, max_seq_len, head_dim]
k = key.transpose(1, 2)
v = value.transpose(1, 2)
# [batch_size, n_local_heads, input_len, max_seq_len]
scores = torch.matmul(q, k.transpose(2, 3))
scores = scores + mask
scores = F.softmax(scores.float(), dim=-1).type_as(q)
# [batch_size, n_local_heads, input_len, head_dim]
output = torch.matmul(scores, v)
# return scores, output.transpose(1, 2).contiguous()
# [batch_size, input_len, hidden_dim]
output = output.transpose(1, 2).contiguous().view(batch_size, input_len, -1)
output = self.o_proj(output)
return scores, output
class TimesFMDecoderLayer(nn.Module):
"""Transformer layer."""
def __init__(
self,
hidden_size: int,
intermediate_size: int,
num_heads: int,
num_kv_heads: int,
head_dim: int,
rms_norm_eps: float = 1e-6,
):
super().__init__()
self.self_attn = TimesFMAttention(
hidden_size=hidden_size,
num_heads=num_heads,
num_kv_heads=num_kv_heads,
head_dim=head_dim,
)
self.mlp = TransformerMLP(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
)
self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
def forward(
self,
hidden_states: torch.Tensor,
mask: torch.Tensor,
paddings: torch.Tensor,
kv_write_indices: torch.Tensor | None = None,
kv_cache: Tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
# Self Attention
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
scores, hidden_states = self.self_attn(
hidden_states=hidden_states,
mask=mask,
kv_write_indices=kv_write_indices,
kv_cache=kv_cache,
)
hidden_states = residual + hidden_states
# MLP
hidden_states = self.mlp(hidden_states, paddings=paddings)
return scores, hidden_states
class StackedDecoder(nn.Module):
"""Stacked transformer layer."""
def __init__(
self,
hidden_size: int,
intermediate_size: int,
num_heads: int,
num_kv_heads: int,
head_dim: int,
num_layers: int,
rms_norm_eps: float = 1e-6,
):
super().__init__()
self.layers = nn.ModuleList()
for _ in range(num_layers):
self.layers.append(
TimesFMDecoderLayer(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
num_heads=num_heads,
num_kv_heads=num_kv_heads,
head_dim=head_dim,
rms_norm_eps=rms_norm_eps,
))
def forward(
self,
hidden_states: torch.Tensor,
paddings: torch.Tensor,
kv_write_indices: torch.Tensor | None = None,
kv_caches: List[Tuple[torch.Tensor, torch.Tensor]] | None = None,
) -> torch.Tensor:
padding_mask = convert_paddings_to_mask(paddings, hidden_states.dtype)
atten_mask = causal_mask(hidden_states)
mask = merge_masks(padding_mask, atten_mask)
for i in range(len(self.layers)):
layer = self.layers[i]
kv_cache = kv_caches[i] if kv_caches is not None else None
_, hidden_states = layer(
hidden_states=hidden_states,
mask=mask,
paddings=paddings,
kv_write_indices=kv_write_indices,
kv_cache=kv_cache,
)
return hidden_states
class PositionalEmbedding(torch.nn.Module):
"""Generates position embedding for a given 1-d sequence.
Attributes:
min_timescale: Start of the geometric index. Determines the periodicity of
the added signal.
max_timescale: End of the geometric index. Determines the frequency of the
added signal.
embedding_dims: Dimension of the embedding to be generated.
"""
def __init__(
self,
embedding_dims: int,
min_timescale: int = 1,
max_timescale: int = 10_000,
) -> None:
super().__init__()
self.min_timescale = min_timescale
self.max_timescale = max_timescale
self.embedding_dims = embedding_dims
def forward(self, seq_length=None, position=None):
"""Generates a Tensor of sinusoids with different frequencies.
Args:
seq_length: an optional Python int defining the output sequence length.
if the `position` argument is specified.
position: [B, seq_length], optional position for each token in the
sequence, only required when the sequence is packed.
Returns:
[B, seqlen, D] if `position` is specified, else [1, seqlen, D]
"""
if position is None:
assert seq_length is not None
# [1, seqlen]
position = torch.arange(seq_length, dtype=torch.float32).unsqueeze(0)
else:
assert position.ndim == 2, position.shape
num_timescales = self.embedding_dims // 2
log_timescale_increment = math.log(
float(self.max_timescale) / float(self.min_timescale)) / max(
num_timescales - 1, 1)
inv_timescales = self.min_timescale * torch.exp(
torch.arange(num_timescales, dtype=torch.float32) *
-log_timescale_increment)
scaled_time = position.unsqueeze(2) * inv_timescales.unsqueeze(0).unsqueeze(
0)
signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2)
# Padding to ensure correct embedding dimension
signal = F.pad(signal, (0, 0, 0, self.embedding_dims % 2))
return signal
class PatchedTimeSeriesDecoder(nn.Module):
"""Patched time-series decoder."""
def __init__(self, config: TimesFMConfig):
super().__init__()
self.config = config
self.input_ff_layer = ResidualBlock(
input_dims=2 * config.patch_len,
output_dims=config.hidden_size,
hidden_dims=config.intermediate_size,
)
self.freq_emb = nn.Embedding(num_embeddings=3,
embedding_dim=config.hidden_size)
self.horizon_ff_layer = ResidualBlock(
input_dims=config.hidden_size,
output_dims=config.horizon_len * (1 + len(config.quantiles)),
hidden_dims=config.intermediate_size,
)
self.stacked_transformer = StackedDecoder(
hidden_size=self.config.hidden_size,
intermediate_size=self.config.intermediate_size,
num_heads=self.config.num_heads,
num_kv_heads=self.config.num_kv_heads,
head_dim=self.config.head_dim,
num_layers=self.config.num_layers,
rms_norm_eps=self.config.rms_norm_eps,
)
if self.config.use_positional_embedding:
self.position_emb = PositionalEmbedding(self.config.hidden_size)
def _forward_transform(
self, inputs: torch.Tensor, patched_pads: torch.Tensor
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
"""Input is of shape [B, N, P]."""
mu, sigma = _masked_mean_std(inputs, patched_pads)
sigma = torch.where(
sigma < self.config.tolerance,
torch.tensor(1.0, dtype=sigma.dtype, device=sigma.device),
sigma,
)
# Normalize each patch
outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
outputs = torch.where(
torch.abs(inputs - self.config.pad_val) < self.config.tolerance,
torch.tensor(self.config.pad_val,
dtype=outputs.dtype,
device=outputs.device),
outputs,
)
return outputs, (mu, sigma)
def _reverse_transform(
self, outputs: torch.Tensor, stats: tuple[torch.Tensor,
torch.Tensor]) -> torch.Tensor:
"""Output is of shape [B, N, P, Q]."""
mu, sigma = stats
return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
def _preprocess_input(
self,
input_ts: torch.Tensor,
input_padding: torch.Tensor,
) -> tuple[
torch.Tensor,
torch.Tensor,
tuple[torch.Tensor, torch.Tensor] | None,
torch.Tensor,
]:
"""Preprocess input for stacked transformer."""
# Reshape into patches (using view for efficiency)
bsize = input_ts.shape[0]
patched_inputs = input_ts.view(bsize, -1, self.config.patch_len)
patched_pads = input_padding.view(bsize, -1, self.config.patch_len)
patched_inputs = torch.where(
torch.abs(patched_pads - 1.0) < self.config.tolerance,
torch.tensor(0.0,
dtype=patched_inputs.dtype,
device=patched_inputs.device),
patched_inputs,
)
patched_pads = torch.where(
torch.abs(patched_inputs - self.config.pad_val) < self.config.tolerance,
torch.tensor(1.0, dtype=patched_pads.dtype, device=patched_pads.device),
patched_pads,
)
patched_inputs, stats = self._forward_transform(patched_inputs,
patched_pads)
# B x N x D
patched_inputs = patched_inputs * (1.0 - patched_pads)
concat_inputs = torch.cat([patched_inputs, patched_pads], dim=-1)
model_input = self.input_ff_layer(concat_inputs)
# A patch should not be padded even if there is at least one zero.
patched_padding = torch.min(patched_pads,
dim=-1)[0] # Get the values from the min result
if self.config.use_positional_embedding:
pos_emb = self.position_emb(model_input.shape[1]).to(model_input.device)
pos_emb = torch.concat([pos_emb] * model_input.shape[0], dim=0)
pos_emb = _shift_padded_seq(patched_padding, pos_emb)
model_input += pos_emb
return model_input, patched_padding, stats, patched_inputs
def _postprocess_output(
self,
model_output: torch.Tensor,
num_outputs: int,
stats: tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
"""Postprocess output of stacked transformer."""
# B x N x (H.Q)
output_ts = self.horizon_ff_layer(model_output)
# Reshape using view
b, n, _ = output_ts.shape
output_ts = output_ts.view(b, n, self.config.horizon_len, num_outputs)
return self._reverse_transform(output_ts, stats)
def forward(
self,
input_ts: torch.Tensor,
input_padding: torch.LongTensor,
freq: torch.Tensor,
) -> torch.Tensor:
num_outputs = len(self.config.quantiles) + 1
model_input, patched_padding, stats, _ = self._preprocess_input(
input_ts=input_ts,
input_padding=input_padding,
)
f_emb = self.freq_emb(freq) # B x 1 x D
model_input += f_emb
model_output = self.stacked_transformer(model_input, patched_padding)
output_ts = self._postprocess_output(model_output, num_outputs, stats)
return output_ts
def decode(
self,
input_ts: torch.Tensor,
paddings: torch.Tensor,
freq: torch.LongTensor,
horizon_len: int,
output_patch_len: int | None = None,
max_len: int | None = None,
return_forecast_on_context: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Auto-regressive decoding without caching.
Args:
input_ts: input time-series and paddings. Time-series shape B x C.
paddings: padding shape B x (C + H) where H is the prediction length.
freq: frequency shape B x 1
horizon_len: prediction length.
output_patch_len: output length to be fetched from one step of
auto-regressive decoding.
max_len: maximum training context length.
return_forecast_on_context: whether to return the model forecast on the
context except the first input patch.
Returns:
Tuple of two forecasting results:
- Point (mean) output predictions as a tensor with shape B x H'.
- Full predictions (mean and quantiles) as a tensor with shape
B x H' x (1 + # quantiles).
In particular, if return_forecast_on_context is True, H' is H plus
the forecastable context length, i.e. context_len - (first) patch_len.
"""
final_out = input_ts
context_len = final_out.shape[1]
full_outputs = []
if max_len is None:
max_len = context_len
if paddings.shape[1] != final_out.shape[1] + horizon_len:
raise ValueError(
"Length of paddings must match length of input + horizon_len:"
f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}")
if output_patch_len is None:
output_patch_len = self.config.horizon_len
num_decode_patches = (horizon_len + output_patch_len -
1) // output_patch_len
for step_index in range(num_decode_patches):
current_padding = paddings[:, 0:final_out.shape[1]]
input_ts = final_out[:, -max_len:]
input_padding = current_padding[:, -max_len:]
fprop_outputs = self(input_ts, input_padding, freq)
if return_forecast_on_context and step_index == 0:
# For the first decodings step, collect the model forecast on the
# context except the unavailable first input batch forecast.
new_full_ts = fprop_outputs[:, 0:-1, 0:self.config.patch_len, :]
new_full_ts = new_full_ts.reshape(new_full_ts.size(0), -1,
new_full_ts.size(3))
full_outputs.append(new_full_ts)
# (full batch, last patch, output_patch_len, index of mean forecast = 0)
new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
new_full_ts = fprop_outputs[:, -1, :output_patch_len, :]
# (full batch, last patch, output_patch_len, all output indices)
full_outputs.append(new_full_ts)
final_out = torch.concatenate([final_out, new_ts], axis=-1)
if return_forecast_on_context:
# `full_outputs` indexing starts at after the first input patch.
full_outputs = torch.concatenate(
full_outputs,
axis=1)[:, :(context_len - self.config.patch_len + horizon_len), :]
else:
# `full_outputs` indexing starts at the forecast horizon.
full_outputs = torch.concatenate(full_outputs, axis=1)[:,
0:horizon_len, :]
return (full_outputs[:, :, 0], full_outputs)
v1/src/timesfm/time_features.py
+215
View File
@@ -0,0 +1,215 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Directory to extract time covariates.
Extract time covariates from datetime.
"""
import numpy as np
import pandas as pd
from pandas.tseries.holiday import EasterMonday
from pandas.tseries.holiday import GoodFriday
from pandas.tseries.holiday import Holiday
from pandas.tseries.holiday import SU
from pandas.tseries.holiday import TH
from pandas.tseries.holiday import USColumbusDay
from pandas.tseries.holiday import USLaborDay
from pandas.tseries.holiday import USMartinLutherKingJr
from pandas.tseries.holiday import USMemorialDay
from pandas.tseries.holiday import USPresidentsDay
from pandas.tseries.holiday import USThanksgivingDay
from pandas.tseries.offsets import DateOffset
from pandas.tseries.offsets import Day
from pandas.tseries.offsets import Easter
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm
# This is 183 to cover half a year (in both directions), also for leap years
# + 17 as Eastern can be between March, 22 - April, 25
MAX_WINDOW = 183 + 17
def _distance_to_holiday(holiday):
"""Return distance to given holiday."""
def _distance_to_day(index):
holiday_date = holiday.dates(
index - pd.Timedelta(days=MAX_WINDOW),
index + pd.Timedelta(days=MAX_WINDOW),
)
assert (
len(holiday_date) != 0 # pylint: disable=g-explicit-length-test
), f"No closest holiday for the date index {index} found."
# It sometimes returns two dates if it is exactly half a year after the
# holiday. In this case, the smaller distance (182 days) is returned.
return (index - holiday_date[0]).days
return _distance_to_day
EasterSunday = Holiday(
"Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)]
)
NewYearsDay = Holiday("New Years Day", month=1, day=1)
SuperBowl = Holiday(
"Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1))
)
MothersDay = Holiday(
"Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2))
)
IndependenceDay = Holiday("Independence Day", month=7, day=4)
ChristmasEve = Holiday("Christmas", month=12, day=24)
ChristmasDay = Holiday("Christmas", month=12, day=25)
NewYearsEve = Holiday("New Years Eve", month=12, day=31)
BlackFriday = Holiday(
"Black Friday",
month=11,
day=1,
offset=[pd.DateOffset(weekday=TH(4)), Day(1)],
)
CyberMonday = Holiday(
"Cyber Monday",
month=11,
day=1,
offset=[pd.DateOffset(weekday=TH(4)), Day(4)],
)
HOLIDAYS = [
EasterMonday,
GoodFriday,
USColumbusDay,
USLaborDay,
USMartinLutherKingJr,
USMemorialDay,
USPresidentsDay,
USThanksgivingDay,
EasterSunday,
NewYearsDay,
SuperBowl,
MothersDay,
IndependenceDay,
ChristmasEve,
ChristmasDay,
NewYearsEve,
BlackFriday,
CyberMonday,
]
class TimeCovariates(object):
"""Extract all time covariates except for holidays."""
def __init__(
self,
datetimes,
normalized=True,
holiday=False,
):
"""Init function.
Args:
datetimes: pandas DatetimeIndex (lowest granularity supported is min)
normalized: whether to normalize features or not
holiday: fetch holiday features or not
Returns:
None
"""
self.normalized = normalized
self.dti = datetimes
self.holiday = holiday
def _minute_of_hour(self):
minutes = np.array(self.dti.minute, dtype=np.float32)
if self.normalized:
minutes = minutes / 59.0 - 0.5
return minutes
def _hour_of_day(self):
hours = np.array(self.dti.hour, dtype=np.float32)
if self.normalized:
hours = hours / 23.0 - 0.5
return hours
def _day_of_week(self):
day_week = np.array(self.dti.dayofweek, dtype=np.float32)
if self.normalized:
day_week = day_week / 6.0 - 0.5
return day_week
def _day_of_month(self):
day_month = np.array(self.dti.day, dtype=np.float32)
if self.normalized:
day_month = day_month / 30.0 - 0.5
return day_month
def _day_of_year(self):
day_year = np.array(self.dti.dayofyear, dtype=np.float32)
if self.normalized:
day_year = day_year / 364.0 - 0.5
return day_year
def _month_of_year(self):
month_year = np.array(self.dti.month, dtype=np.float32)
if self.normalized:
month_year = month_year / 11.0 - 0.5
return month_year
def _week_of_year(self):
week_year = np.array(self.dti.strftime("%U").astype(int), dtype=np.float32)
if self.normalized:
week_year = week_year / 51.0 - 0.5
return week_year
def _get_holidays(self):
dti_series = self.dti.to_series()
hol_variates = np.vstack([
dti_series.apply(_distance_to_holiday(h)).values for h in tqdm(HOLIDAYS)
])
# hol_variates is (num_holiday, num_time_steps), the normalization should be
# performed in the num_time_steps dimension.
return StandardScaler().fit_transform(hol_variates.T).T
def get_covariates(self):
"""Get all time covariates."""
moh = self._minute_of_hour().reshape(1, -1)
hod = self._hour_of_day().reshape(1, -1)
dom = self._day_of_month().reshape(1, -1)
dow = self._day_of_week().reshape(1, -1)
doy = self._day_of_year().reshape(1, -1)
moy = self._month_of_year().reshape(1, -1)
woy = self._week_of_year().reshape(1, -1)
all_covs = [
moh,
hod,
dom,
dow,
doy,
moy,
woy,
]
columns = ["moh", "hod", "dom", "dow", "doy", "moy", "woy"]
if self.holiday:
hol_covs = self._get_holidays()
all_covs.append(hol_covs)
columns += [f"hol_{i}" for i in range(len(HOLIDAYS))]
return pd.DataFrame(
data=np.vstack(all_covs).transpose(),
columns=columns,
index=self.dti,
)
v1/src/timesfm/timesfm_base.py
+736
View File
@@ -0,0 +1,736 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Base class for TimesFM inference. This will be common to PAX and Pytorch."""
import collections
import dataclasses
import logging
import multiprocessing
from typing import Any, Literal, Sequence, TYPE_CHECKING
import numpy as np
import pandas as pd
from utilsforecast.processing import make_future_dataframe
if TYPE_CHECKING:
from . import xreg_lib
Category = xreg_lib.Category
XRegMode = xreg_lib.XRegMode
else:
Category = int | str
XRegMode = str
_TOL = 1e-6
DEFAULT_QUANTILES = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
def process_group(key, group, value_name, forecast_context_len):
group = group.tail(forecast_context_len)
return np.array(group[value_name], dtype=np.float32), key
def moving_average(arr, window_size):
"""Calculates the moving average using NumPy's convolution function."""
# Pad with zeros to handle initial window positions
arr_padded = np.pad(arr, (window_size - 1, 0), "constant")
smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") /
window_size)
return [smoothed_arr, arr - smoothed_arr]
def freq_map(freq: str):
"""Returns the frequency map for the given frequency string."""
freq = str.upper(freq)
if freq.endswith("MS"):
return 1
elif freq.endswith(("H", "T", "MIN", "D", "B", "U", "S")):
return 0
elif (
freq.endswith(("W", "M"))
or freq.startswith("W-")
or (freq.startswith("M") and len(freq) == 2)
):
return 1
elif (
freq.endswith(("Y", "Q", "A"))
or freq.startswith("Y-")
or freq.startswith("Q-")
or freq.startswith("A-")
):
return 2
else:
raise ValueError(f"Invalid frequency: {freq}")
def strip_leading_nans(arr):
"""
Removes contiguous NaN values from the beginning of a NumPy array.
Args:
arr: The input NumPy array.
Returns:
A new NumPy array with leading NaN values removed.
If the array is all NaNs or empty, returns an empty array.
"""
isnan = np.isnan(arr)
first_valid_index = np.argmax(~isnan)
return arr[first_valid_index:]
def linear_interpolation(arr):
"""
Performs linear interpolation to fill NaN values in a 1D numpy array.
Args:
arr: The 1D numpy array containing NaN values.
Returns:
A new numpy array with NaN values filled using linear interpolation,
or the original array if no NaNs are present.
Returns None if the input is not a 1D array.
Returns the original array if there are no NaN values.
"""
nans = np.isnan(arr)
if not np.any(nans): # Check if there are any NaNs
return arr
def x(z):
return z.nonzero()[0]
nans_indices = x(nans)
non_nans_indices = x(~nans)
non_nans_values = arr[~nans]
try:
arr[nans] = np.interp(nans_indices, non_nans_indices, non_nans_values)
except ValueError:
if len(non_nans_values) > 0:
mu = np.nanmean(arr)
else:
mu = 0.0
arr = np.where(np.isfinite(arr), arr, mu)
return arr
# Per time series normalization: forward.
def _normalize(batch):
stats = [
(np.mean(x), np.where((w := np.std(x)) > _TOL, w, 1.0)) for x in batch
]
new_batch = [(x - stat[0]) / stat[1] for x, stat in zip(batch, stats)]
return new_batch, stats
# Per time series normalization: inverse.
def _renormalize(batch, stats):
return [x * stat[1] + stat[0] for x, stat in zip(batch, stats)]
@dataclasses.dataclass(kw_only=True)
class TimesFmHparams:
"""Hparams used to initialize a TimesFM model for inference.
These are the sufficient subset of hparams to configure TimesFM inference
agnostic to the checkpoint version, and are not necessarily the same as the
hparams used to train the checkpoint.
Attributes:
context_len: Largest context length the model allows for each decode call.
This technically can be any large, but practically should set to the
context length the checkpoint was trained with.
horizon_len: Forecast horizon.
input_patch_len: Input patch len.
output_patch_len: Output patch len. How many timepoints is taken from a
single step of autoregressive decoding. Can be set as the training horizon
of the checkpoint.
num_layers: Number of transformer layers in the model.
model_dims: Model dimension.
per_core_batch_size: Batch size on each core for data parallelism.
backend: One of "cpu", "gpu" or "tpu".
quantiles: Which quantiles are output by the model.
"""
context_len: int = 512
horizon_len: int = 128
input_patch_len: int = 32
output_patch_len: int = 128
num_layers: int = 20
num_heads: int = 16
model_dims: int = 1280
per_core_batch_size: int = 32
backend: Literal["cpu", "gpu", "tpu"] = "cpu"
quantiles: Sequence[float] | None = DEFAULT_QUANTILES
use_positional_embedding: bool = True
# Hparams beyond the model.
point_forecast_mode: Literal["mean", "median"] = "median"
@dataclasses.dataclass(kw_only=True)
class TimesFmCheckpoint:
"""Checkpoint used to initialize a TimesFM model for inference.
Attributes:
version: Version of the checkpoint, e.g. "jax", "torch", "tensorflow", etc.
The factory will create the corresponding TimesFm inference class based on
this version.
path: Path to the checkpoint.
type: If provided, type of the checkpoint used by the specific checkpoint
loader per version.
step: If provided, step of the checkpoint.
"""
version: str = "jax"
path: str | None = None
huggingface_repo_id: str | None = None
type: Any = None
step: int | None = None
local_dir: str | None = None
class TimesFmBase:
"""Base TimesFM forecast API for inference.
This class is the scaffolding for calling TimesFM forecast. To properly use:
1. Create an instance with the correct hyperparameters of a TimesFM model.
2. Call `load_from_checkpoint` to load a compatible checkpoint.
3. Call `forecast` for inference.
"""
def _logging(self, s):
print(s)
def __post_init__(self) -> None:
"""Additional initialization for subclasses before checkpoint loading."""
pass
def __init__(self, hparams: TimesFmHparams,
checkpoint: TimesFmCheckpoint) -> None:
"""Initializes the TimesFM forecast API.
Args:
hparams: Hyperparameters of the model.
checkpoint: Checkpoint to load. Notice `checkpoint.version` will decide
which TimesFM version to use.
"""
self.hparams = hparams
# Expand hparams for conciseness within the model code.
self.context_len = hparams.context_len
self.horizon_len = hparams.horizon_len
self.input_patch_len = hparams.input_patch_len
self.output_patch_len = hparams.output_patch_len
self.num_layers = hparams.num_layers
self.model_dims = hparams.model_dims
self.backend = hparams.backend
self.quantiles = hparams.quantiles
self.num_heads = hparams.num_heads
self.use_pos_emb = hparams.use_positional_embedding
# Rewrite these values in __post_init__ for SPMD.
self.num_cores = 1
self.per_core_batch_size = hparams.per_core_batch_size
self.global_batch_size = hparams.per_core_batch_size
self._horizon_start = self.context_len - self.input_patch_len
self.__post_init__()
self.load_from_checkpoint(checkpoint)
def load_from_checkpoint(self, checkpoint: TimesFmCheckpoint) -> None:
"""Loads a checkpoint and compiles the decoder."""
raise NotImplementedError("`load_from_checkpoint` is not implemented.")
def _preprocess(
self, inputs: Sequence[np.ndarray],
freq: Sequence[int]) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
"""Formats and pads raw inputs to feed into the model.
This function both pads each time series to match the context length, and
pads the inputs to meet the SPMD shape requirement.
Args:
inputs: A list of 1d JTensors. Each JTensor is the context time series of
a single forecast task.
freq: list of frequencies
Returns:
A tuple of:
- the padded input time series to meet the model required context.
- the padding indicator.
- the frequency of each input time series.
- the number of padded examples for SPMD so that each core has the same
number (a multiple of `batch_size`) of examples.
"""
input_ts, input_padding, inp_freq = [], [], []
pmap_pad = ((len(inputs) - 1) // self.global_batch_size +
1) * self.global_batch_size - len(inputs)
for i, ts in enumerate(inputs):
input_len = ts.shape[0]
padding = np.zeros(shape=(input_len + self.horizon_len,), dtype=float)
if input_len < self.context_len:
num_front_pad = self.context_len - input_len
ts = np.concatenate([np.zeros(shape=(num_front_pad,), dtype=float), ts],
axis=0)
padding = np.concatenate(
[np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0)
elif input_len > self.context_len:
ts = ts[-self.context_len:]
padding = padding[-(self.context_len + self.horizon_len):]
input_ts.append(ts)
input_padding.append(padding)
inp_freq.append(freq[i])
# Padding the remainder batch.
for _ in range(pmap_pad):
input_ts.append(input_ts[-1])
input_padding.append(input_padding[-1])
inp_freq.append(inp_freq[-1])
return (
np.stack(input_ts, axis=0),
np.stack(input_padding, axis=0),
np.array(inp_freq).astype(np.int32).reshape(-1, 1),
pmap_pad,
)
def _forecast(
self,
inputs: Sequence[Any],
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts on a list of time series.
Args:
inputs: list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
freq: frequency of each context time series. 0 for high frequency
(default), 1 for medium, and 2 for low. Notice this is different from
the `freq` required by `forecast_on_df`.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
forecast_context_len: optional max context length.
return_forecast_on_context: True to return the forecast on the context
when available, i.e. after the first input patch.
Returns:
A tuple for np.array:
- the mean forecast of size (# inputs, # forecast horizon),
- the full forecast (mean + quantiles) of size
(# inputs, # forecast horizon, 1 + # quantiles).
Raises:
ValueError: If the checkpoint is not properly loaded.
"""
raise NotImplementedError("`_forecast` is not implemented.")
def forecast(
self,
inputs: Sequence[Any],
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
normalize: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts on a list of time series.
Args:
inputs: list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
freq: frequency of each context time series. 0 for high frequency
(default), 1 for medium, and 2 for low. Notice this is different from
the `freq` required by `forecast_on_df`.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
forecast_context_len: optional max context length.
return_forecast_on_context: True to return the forecast on the context
when available, i.e. after the first input patch.
normalize: If True, then we normalize the inputs before forecasting and
the outputs are then renormalized to the original scale.
Returns:
A tuple for np.array:
- the mean forecast of size (# inputs, # forecast horizon),
- the full forecast (mean + quantiles) of size
(# inputs, # forecast horizon, 1 + # quantiles).
Raises:
ValueError: If the checkpoint is not properly loaded.
"""
stats = None
tmp_inputs = []
for each_input in inputs:
arr = np.array(each_input)
if not np.isfinite(arr).all():
arr = np.where(np.isfinite(arr), arr, np.nan)
arr = strip_leading_nans(arr)
arr = linear_interpolation(arr)
tmp_inputs.append(arr)
inputs = tmp_inputs
if normalize:
inputs, stats = _normalize(inputs)
mean_forecast, quantile_forecast = self._forecast(
inputs,
freq,
window_size,
forecast_context_len,
return_forecast_on_context,
)
if stats is not None:
stats = np.array(stats)
mu = stats[:, 0]
sigma = stats[:, 1]
mean_forecast = mean_forecast * sigma[:, None] + mu[:, None]
quantile_forecast = (quantile_forecast * sigma[:, None, None] +
mu[:, None, None])
if self.hparams.point_forecast_mode == "mean":
return mean_forecast, quantile_forecast
elif self.hparams.point_forecast_mode == "median":
if self._median_index == -1:
for i, quantile in enumerate(self.quantiles):
if quantile == 0.5:
self._median_index = i
break
if self._median_index == -1:
raise ValueError("Median (0.5) is not found in the model quantiles:"
f" {self.quantiles}. Please check the hparams.")
return (
quantile_forecast[:, :, 1 + self._median_index],
quantile_forecast,
)
else:
raise ValueError(
"Unsupported point forecast mode:"
f" {self.hparams.point_forecast_mode}. Use 'mean' or 'median'.")
def forecast_with_covariates(
self,
inputs: list[Sequence[float]],
dynamic_numerical_covariates: (dict[str, Sequence[Sequence[float]]] |
None) = None,
dynamic_categorical_covariates: (dict[str, Sequence[Sequence[Category]]] |
None) = None,
static_numerical_covariates: dict[str, Sequence[float]] | None = None,
static_categorical_covariates: (dict[str, Sequence[Category]] |
None) = None,
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
xreg_mode: XRegMode = "xreg + timesfm",
normalize_xreg_target_per_input: bool = True,
ridge: float = 0.0,
max_rows_per_col: int = 0,
force_on_cpu: bool = False,
):
"""Forecasts on a list of time series with covariates.
To optimize inference speed, avoid string valued categorical covariates.
Args:
inputs: A list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
dynamic_numerical_covariates: A dict of dynamic numerical covariates.
dynamic_categorical_covariates: A dict of dynamic categorical covariates.
static_numerical_covariates: A dict of static numerical covariates.
static_categorical_covariates: A dict of static categorical covariates.
freq: frequency of each context time series. 0 for high frequency
(default), 1 for medium, and 2 for low. Notice this is different from
the `freq` required by `forecast_on_df`.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
forecast_context_len: optional max context length.
xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "xreg + timesfm"
fits a model on the residuals of the TimesFM forecast. "timesfm + xreg"
fits a model on the targets then forecasts on the residuals via TimesFM.
normalize_xreg_target_per_input: whether to normalize the xreg target per
input in the given batch.
ridge: ridge penalty for the linear model.
max_rows_per_col: max number of rows per column for the linear model.
force_on_cpu: whether to force running on cpu for the linear model.
Returns:
A tuple of two lists. The first is the outputs of the model. The second is
the outputs of the xreg.
"""
from . import xreg_lib
# Verify and bookkeep covariates.
if not (dynamic_numerical_covariates or dynamic_categorical_covariates or
static_numerical_covariates or static_categorical_covariates):
raise ValueError(
"At least one of dynamic_numerical_covariates,"
" dynamic_categorical_covariates, static_numerical_covariates,"
" static_categorical_covariates must be set.")
# Track the lengths of (1) each input, (2) the part that can be used in the
# linear model, and (3) the horizon.
input_lens, train_lens, test_lens = [], [], []
for i, input_ts in enumerate(inputs):
input_len = len(input_ts)
input_lens.append(input_len)
if xreg_mode == "timesfm + xreg":
# For fitting residuals, no TimesFM forecast on the first patch.
train_lens.append(max(0, input_len - self.input_patch_len))
elif xreg_mode == "xreg + timesfm":
train_lens.append(input_len)
else:
raise ValueError(f"Unsupported mode: {xreg_mode}")
if dynamic_numerical_covariates:
test_lens.append(
len(list(dynamic_numerical_covariates.values())[0][i]) - input_len)
elif dynamic_categorical_covariates:
test_lens.append(
len(list(dynamic_categorical_covariates.values())[0][i]) -
input_len)
else:
test_lens.append(self.horizon_len)
if test_lens[-1] > self.horizon_len:
raise ValueError(
"Forecast requested longer horizon than the model definition "
f"supports: {test_lens[-1]} vs {self.horizon_len}.")
# Prepare the covariates into train and test.
train_dynamic_numerical_covariates = collections.defaultdict(list)
test_dynamic_numerical_covariates = collections.defaultdict(list)
train_dynamic_categorical_covariates = collections.defaultdict(list)
test_dynamic_categorical_covariates = collections.defaultdict(list)
for covariates, train_covariates, test_covariates in (
(
dynamic_numerical_covariates,
train_dynamic_numerical_covariates,
test_dynamic_numerical_covariates,
),
(
dynamic_categorical_covariates,
train_dynamic_categorical_covariates,
test_dynamic_categorical_covariates,
),
):
if not covariates:
continue
for covariate_name, covariate_values in covariates.items():
for input_len, train_len, covariate_value in zip(
input_lens, train_lens, covariate_values):
train_covariates[covariate_name].append(
covariate_value[(input_len - train_len):input_len])
test_covariates[covariate_name].append(covariate_value[input_len:])
# Fit models.
if xreg_mode == "timesfm + xreg":
# Forecast via TimesFM then fit a model on the residuals.
mean_outputs, _ = self.forecast(
inputs,
freq,
window_size,
forecast_context_len,
return_forecast_on_context=True,
)
targets = [
(np.array(input_ts)[-train_len:] -
mean_output[(self._horizon_start - train_len):self._horizon_start])
for input_ts, mean_output, train_len in zip(inputs, mean_outputs,
train_lens)
]
per_instance_stats = None
if normalize_xreg_target_per_input:
targets, per_instance_stats = _normalize(targets)
xregs = xreg_lib.BatchedInContextXRegLinear(
targets=targets,
train_lens=train_lens,
test_lens=test_lens,
train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
train_dynamic_categorical_covariates=
train_dynamic_categorical_covariates,
test_dynamic_categorical_covariates=
test_dynamic_categorical_covariates,
static_numerical_covariates=static_numerical_covariates,
static_categorical_covariates=static_categorical_covariates,
).fit(
ridge=ridge,
one_hot_encoder_drop=None if ridge > 0 else "first",
max_rows_per_col=max_rows_per_col,
force_on_cpu=force_on_cpu,
debug_info=False,
assert_covariates=True,
assert_covariate_shapes=True,
)
if normalize_xreg_target_per_input:
xregs = _renormalize(xregs, per_instance_stats)
outputs = [
(mean_output[self._horizon_start:(self._horizon_start + test_len)] +
xreg)
for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
]
else:
# Fit a model on the targets then forecast on the residuals via TimesFM.
targets = [
np.array(input_ts)[-train_len:]
for input_ts, train_len in zip(inputs, train_lens)
]
per_instance_stats = None
if normalize_xreg_target_per_input:
targets, per_instance_stats = _normalize(targets)
xregs, xregs_on_context, _, _, _ = xreg_lib.BatchedInContextXRegLinear(
targets=targets,
train_lens=train_lens,
test_lens=test_lens,
train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
train_dynamic_categorical_covariates=
train_dynamic_categorical_covariates,
test_dynamic_categorical_covariates=
test_dynamic_categorical_covariates,
static_numerical_covariates=static_numerical_covariates,
static_categorical_covariates=static_categorical_covariates,
).fit(
ridge=ridge,
one_hot_encoder_drop=None if ridge > 0 else "first",
max_rows_per_col=max_rows_per_col,
force_on_cpu=force_on_cpu,
debug_info=True,
assert_covariates=True,
assert_covariate_shapes=True,
)
mean_outputs, _ = self.forecast(
[
target - xreg_on_context
for target, xreg_on_context in zip(targets, xregs_on_context)
],
freq,
window_size,
forecast_context_len,
return_forecast_on_context=True,
)
outputs = [
(mean_output[self._horizon_start:(self._horizon_start + test_len)] +
xreg)
for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
]
if normalize_xreg_target_per_input:
outputs = _renormalize(outputs, per_instance_stats)
return outputs, xregs
def forecast_on_df(
self,
inputs: pd.DataFrame,
freq: str,
forecast_context_len: int = 0,
value_name: str = "values",
model_name: str = "timesfm",
window_size: int | None = None,
num_jobs: int = 1,
normalize: bool = False,
verbose: bool = True,
) -> pd.DataFrame:
"""Forecasts on a list of time series.
Args:
inputs: A pd.DataFrame of all time series. The dataframe should have a
`unique_id` column for identifying the time series, a `ds` column for
timestamps and a value column for the time series values.
freq: string valued `freq` of data. Notice this is different from the
`freq` required by `forecast`. See `freq_map` for allowed values.
forecast_context_len: If provided none zero, we take the last
`forecast_context_len` time-points from each series as the forecast
context instead of the `context_len` set by the model.
value_name: The name of the value column.
model_name: name of the model to be written into future df.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
num_jobs: number of parallel processes to use for dataframe processing.
normalize: normalize context before forecasting or not.
verbose: output model states in terminal.
Returns:
Future forecasts dataframe.
"""
if not ("unique_id" in inputs.columns and "ds" in inputs.columns and
value_name in inputs.columns):
raise ValueError(
f"DataFrame must have unique_id, ds and {value_name} columns.")
if not forecast_context_len:
forecast_context_len = self.context_len
logging.info("Preprocessing dataframe.")
df_sorted = inputs.sort_values(by=["unique_id", "ds"])
new_inputs = []
uids = []
if num_jobs == 1:
if verbose:
print("Processing dataframe with single process.")
for key, group in df_sorted.groupby("unique_id"):
inp, uid = process_group(
key,
group,
value_name,
forecast_context_len,
)
new_inputs.append(inp)
uids.append(uid)
else:
if num_jobs == -1:
num_jobs = multiprocessing.cpu_count()
if verbose:
print("Processing dataframe with multiple processes.")
with multiprocessing.Pool(processes=num_jobs) as pool:
results = pool.starmap(
process_group,
[(key, group, value_name, forecast_context_len)
for key, group in df_sorted.groupby("unique_id")],
)
new_inputs, uids = zip(*results)
if verbose:
print("Finished preprocessing dataframe.")
freq_inps = [freq_map(freq)] * len(new_inputs)
_, full_forecast = self.forecast(new_inputs,
freq=freq_inps,
normalize=normalize,
window_size=window_size)
if verbose:
print("Finished forecasting.")
fcst_df = make_future_dataframe(
uids=uids,
last_times=df_sorted.groupby("unique_id")["ds"].tail(1),
h=self.horizon_len,
freq=freq,
)
fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
for i, q in enumerate(self.quantiles):
q_col = f"{model_name}-q-{q}"
fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
1 + i].reshape(-1, 1)
if q == 0.5:
fcst_df[model_name] = fcst_df[q_col]
logging.info("Finished creating output dataframe.")
return fcst_df
v1/src/timesfm/timesfm_jax.py
+353
View File
@@ -0,0 +1,353 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TimesFM JAX forecast API for inference."""
import logging
import multiprocessing
import time
from os import path
from typing import Any, Sequence
import einshape as es
import jax
import jax.numpy as jnp
import numpy as np
from huggingface_hub import snapshot_download
from paxml import checkpoints, tasks_lib
from praxis import base_hyperparams, base_layer, pax_fiddle, py_utils, pytypes
from praxis.layers import normalizations, transformers
from timesfm import timesfm_base
from timesfm import patched_decoder
instantiate = base_hyperparams.instantiate
NestedMap = py_utils.NestedMap
JTensor = pytypes.JTensor
_TOL = 1e-6
class TimesFmJax(timesfm_base.TimesFmBase):
"""TimesFM forecast API for inference.
This class is the scaffolding for calling TimesFM forecast. To properly use:
1. Create an instance with the correct hyperparameters of a TimesFM model.
2. Call `load_from_checkpoint` to load a compatible checkpoint.
3. Call `forecast` for inference.
Given the model size, this API does not shard the model weights for SPMD. All
parallelism happens on the data dimension.
Compilation happens during the first time `forecast` is called and uses the
`per_core_batch_size` to set and freeze the input signature. Subsequent calls
to `forecast` reflect the actual inference latency.
"""
def _get_sample_inputs(self):
return {
"input_ts":
jnp.zeros(
(
self.per_core_batch_size,
self.context_len + self.output_patch_len,
),
dtype=jnp.float32,
),
"input_padding":
jnp.zeros(
(
self.per_core_batch_size,
self.context_len + self.output_patch_len,
),
dtype=jnp.float32,
),
"freq":
jnp.zeros(
(
self.per_core_batch_size,
1,
),
dtype=jnp.int32,
),
}
def __post_init__(self):
self.num_cores = jax.local_device_count(self.backend)
self.global_batch_size = self.per_core_batch_size * self.num_cores
self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
self._pmapped_decode = None
self._model = None
self._train_state = None
self._median_index = -1
def load_from_checkpoint(
self,
checkpoint: timesfm_base.TimesFmCheckpoint,
) -> None:
"""Loads a checkpoint and compiles the decoder."""
checkpoint_type = (checkpoints.CheckpointType.FLAX
if checkpoint.type is None else checkpoint.type)
checkpoint_path = checkpoint.path
step = checkpoint.step
repo_id = checkpoint.huggingface_repo_id
if checkpoint_path is None:
checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints")
# Rewrite the devices for Jax.
self.mesh_shape = [1, self.num_cores, 1]
self.mesh_name = ["replica", "data", "mdl"]
self.model_p = pax_fiddle.Config(
patched_decoder.PatchedTimeSeriesDecoder,
name="patched_decoder",
horizon_len=self.output_patch_len,
patch_len=self.input_patch_len,
model_dims=self.model_dims,
hidden_dims=self.model_dims,
residual_block_tpl=pax_fiddle.Config(patched_decoder.ResidualBlock),
quantiles=self.quantiles,
use_freq=True,
use_pos_emb=self.use_pos_emb,
stacked_transformer_params_tpl=pax_fiddle.Config(
transformers.StackedTransformer,
num_heads=self.num_heads,
num_layers=self.num_layers,
transformer_layer_params_tpl=pax_fiddle.Config(
transformers.Transformer,
ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
),
),
)
self._key1, self._key2 = jax.random.split(jax.random.PRNGKey(42))
self._model = None
self._train_state = None
self._pmapped_decode = None
self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
try:
multiprocessing.set_start_method("spawn")
except RuntimeError:
print("Multiprocessing context has already been set.")
# Download the checkpoint from Hugging Face Hub if not given
# Initialize the model weights.
self._logging("Constructing model weights.")
start_time = time.time()
self._model = instantiate(self.model_p)
var_weight_hparams = self._model.abstract_init_with_metadata(
self._get_sample_inputs(), do_eval=True)
train_state_partition_specs = tasks_lib.create_state_partition_specs(
var_weight_hparams,
mesh_shape=self.mesh_shape,
mesh_axis_names=self.mesh_name,
discard_opt_states=True,
learners=None,
)
train_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
var_weight_hparams,
discard_opt_states=True,
learners=None,
)
self._logging(
f"Constructed model weights in {time.time() - start_time:.2f} seconds.")
# Load the model weights.
self._logging(f"Restoring checkpoint from {checkpoint_path}.")
start_time = time.time()
self._train_state = checkpoints.restore_checkpoint(
train_state_local_shapes,
checkpoint_dir=checkpoint_path,
checkpoint_type=checkpoint_type,
state_specs=train_state_partition_specs,
step=step,
)
self._logging(
f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
self.jit_decode()
def jit_decode(self):
"""Jitting decoding function."""
# Initialize and jit the decode fn.
def _decode(inputs):
assert self._model is not None
assert self._train_state is not None
return self._model.apply(
self._train_state.mdl_vars,
inputs,
horizon_len=self.horizon_len,
output_patch_len=self.output_patch_len,
max_len=self.context_len,
return_forecast_on_context=True,
rngs={
base_layer.PARAMS: self._key1,
base_layer.RANDOM: self._key2,
},
method=self._model.decode,
)
self._logging("Jitting decoding.")
start_time = time.time()
self._pmapped_decode = jax.pmap(
_decode,
axis_name="batch",
devices=jax.devices(self.backend),
backend=self.backend,
axis_size=self.num_cores,
)
with base_layer.JaxContext.new_context(hparams=self._eval_context):
_ = self._pmapped_decode(
NestedMap({
"input_ts":
jnp.zeros(
(
self.num_cores,
self.per_core_batch_size,
self.context_len,
),
dtype=jnp.float32,
),
"input_padding":
jnp.zeros(
(
self.num_cores,
self.per_core_batch_size,
self.context_len + self.horizon_len,
),
dtype=jnp.float32,
),
"date_features":
None,
"freq":
jnp.zeros(
(self.num_cores, self.per_core_batch_size, 1),
dtype=jnp.int32,
),
}))
self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
def _forecast(
self,
inputs: Sequence[Any],
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts on a list of time series.
Args:
inputs: list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
freq: frequency of each context time series. 0 for high frequency
(default), 1 for medium, and 2 for low. Notice this is different from
the `freq` required by `forecast_on_df`.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
forecast_context_len: optional max context length.
return_forecast_on_context: True to return the forecast on the context
when available, i.e. after the first input patch.
Returns:
A tuple for JTensors:
- the mean forecast of size (# inputs, # forecast horizon),
- the full forecast (mean + quantiles) of size
(# inputs, # forecast horizon, 1 + # quantiles).
Raises:
ValueError: If the checkpoint is not properly loaded.
"""
if not self._train_state or not self._model:
raise ValueError(
"Checkpoint not loaded. Call `load_from_checkpoint` before"
" `forecast`.")
if forecast_context_len is None:
fcontext_len = self.context_len
else:
fcontext_len = forecast_context_len
inputs = [np.array(ts)[-fcontext_len:] for ts in inputs]
if window_size is not None:
new_inputs = []
for ts in inputs:
new_inputs.extend(timesfm_base.moving_average(ts, window_size))
inputs = new_inputs
if freq is None:
logging.info("No frequency provided via `freq`. Default to high (0).")
freq = [0] * len(inputs)
input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
with base_layer.JaxContext.new_context(hparams=self._eval_context):
mean_outputs = []
full_outputs = []
assert input_ts.shape[0] % self.global_batch_size == 0
for i in range(input_ts.shape[0] // self.global_batch_size):
input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
self.global_batch_size])
input_padding_in = jnp.array(
input_padding[i * self.global_batch_size:(i + 1) *
self.global_batch_size],)
inp_freq_in = jnp.array(
inp_freq[i * self.global_batch_size:(i + 1) *
self.global_batch_size, :],
dtype=jnp.int32,
)
pmapped_inputs = NestedMap({
"input_ts":
es.jax_einshape(
"(db)...->db...",
input_ts_in,
d=self.num_cores,
),
"input_padding":
es.jax_einshape(
"(db)...->db...",
input_padding_in,
d=self.num_cores,
),
"date_features":
None,
"freq":
es.jax_einshape(
"(db)...->db...",
inp_freq_in,
d=self.num_cores,
),
})
mean_output, full_output = self._pmapped_decode(pmapped_inputs)
if not return_forecast_on_context:
mean_output = mean_output[:, :, self._horizon_start:, ...]
full_output = full_output[:, :, self._horizon_start:, ...]
mean_output = es.jax_einshape("db...->(db)...",
mean_output,
d=self.num_cores)
full_output = es.jax_einshape("db...->(db)...",
full_output,
d=self.num_cores)
mean_output = np.array(mean_output)
full_output = np.array(full_output)
mean_outputs.append(mean_output)
full_outputs.append(full_output)
mean_outputs = np.concatenate(mean_outputs, axis=0)
full_outputs = np.concatenate(full_outputs, axis=0)
if pmap_pad > 0:
mean_outputs = mean_outputs[:-pmap_pad, ...]
full_outputs = full_outputs[:-pmap_pad, ...]
if window_size is not None:
mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
return mean_outputs, full_outputs
v1/src/timesfm/timesfm_torch.py
+168
View File
@@ -0,0 +1,168 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TimesFM pytorch forecast API for inference."""
import logging
from os import path
from typing import Any, Sequence
import numpy as np
import torch
from huggingface_hub import snapshot_download
from timesfm import timesfm_base
from . import pytorch_patched_decoder as ppd
_TOL = 1e-6
class TimesFmTorch(timesfm_base.TimesFmBase):
"""TimesFM forecast API for inference."""
def __post_init__(self):
self._model_config = ppd.TimesFMConfig(
num_layers=self.num_layers,
num_heads=self.num_heads,
hidden_size=self.model_dims,
intermediate_size=self.model_dims,
patch_len=self.input_patch_len,
horizon_len=self.output_patch_len,
head_dim=self.model_dims // self.num_heads,
quantiles=self.quantiles,
use_positional_embedding=self.use_pos_emb,
)
self._model = None
self.num_cores = 1
self.global_batch_size = self.per_core_batch_size
self._device = torch.device("cuda:0" if (
torch.cuda.is_available() and self.backend == "gpu") else "cpu")
self._median_index = -1
def load_from_checkpoint(
self,
checkpoint: timesfm_base.TimesFmCheckpoint,
) -> None:
"""Loads a checkpoint and compiles the decoder."""
checkpoint_path = checkpoint.path
repo_id = checkpoint.huggingface_repo_id
if checkpoint_path is None:
checkpoint_path = path.join(
snapshot_download(repo_id, local_dir=checkpoint.local_dir),
"torch_model.ckpt")
self._model = ppd.PatchedTimeSeriesDecoder(self._model_config)
loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
logging.info("Loading checkpoint from %s", checkpoint_path)
self._model.load_state_dict(loaded_checkpoint)
logging.info("Sending checkpoint to device %s", f"{self._device}")
self._model.to(self._device)
self._model.eval()
# TODO: add compilation.
def _forecast(
self,
inputs: Sequence[Any],
freq: Sequence[int] | None = None,
window_size: int | None = None,
forecast_context_len: int | None = None,
return_forecast_on_context: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
"""Forecasts on a list of time series.
Args:
inputs: list of time series forecast contexts. Each context time series
should be in a format convertible to JTensor by `jnp.array`.
freq: frequency of each context time series. 0 for high frequency
(default), 1 for medium, and 2 for low. Notice this is different from
the `freq` required by `forecast_on_df`.
window_size: window size of trend + residual decomposition. If None then
we do not do decomposition.
forecast_context_len: optional max context length.
return_forecast_on_context: True to return the forecast on the context
when available, i.e. after the first input patch.
Returns:
A tuple for JTensors:
- the mean forecast of size (# inputs, # forecast horizon),
- the full forecast (mean + quantiles) of size
(# inputs, # forecast horizon, 1 + # quantiles).
Raises:
ValueError: If the checkpoint is not properly loaded.
"""
if self._model is None:
raise ValueError("Checkpoint is not properly loaded.")
if forecast_context_len is None:
forecast_context_len = self.context_len
inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs]
if window_size is not None:
new_inputs = []
for ts in inputs:
new_inputs.extend(timesfm_base.moving_average(ts, window_size))
inputs = new_inputs
if freq is None:
logging.info("No frequency provided via `freq`. Default to high (0).")
freq = [0] * len(inputs)
input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
with torch.no_grad():
mean_outputs = []
full_outputs = []
for i in range(input_ts.shape[0] // self.global_batch_size):
t_input_ts = torch.Tensor(input_ts[i * self.global_batch_size:(i + 1) *
self.global_batch_size]).to(
self._device)
t_input_padding = torch.Tensor(
input_padding[i * self.global_batch_size:(i + 1) *
self.global_batch_size]).to(self._device)
t_inp_freq = torch.LongTensor(
inp_freq[i * self.global_batch_size:(i + 1) *
self.global_batch_size, :]).to(self._device)
mean_output, full_output = self._model.decode(
input_ts=t_input_ts,
paddings=t_input_padding,
freq=t_inp_freq,
horizon_len=self.horizon_len,
output_patch_len=self.output_patch_len,
# Returns forecasts on context for parity with the Jax version.
return_forecast_on_context=True,
)
if not return_forecast_on_context:
mean_output = mean_output[:, self._horizon_start:, ...]
full_output = full_output[:, self._horizon_start:, ...]
if self.backend == "gpu":
mean_output = mean_output.cpu()
full_output = full_output.cpu()
mean_output = mean_output.detach().numpy()
full_output = full_output.detach().numpy()
mean_outputs.append(mean_output)
full_outputs.append(full_output)
mean_outputs = np.concatenate(mean_outputs, axis=0)
full_outputs = np.concatenate(full_outputs, axis=0)
if pmap_pad > 0:
mean_outputs = mean_outputs[:-pmap_pad, ...]
full_outputs = full_outputs[:-pmap_pad, ...]
if window_size is not None:
mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
return mean_outputs, full_outputs
v1/src/timesfm/xreg_lib.py
+486
View File
@@ -0,0 +1,486 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Helper functions for in-context covariates and regression."""
import itertools
import math
from typing import Any, Iterable, Literal, Mapping, Sequence
import jax
import jax.numpy as jnp
import numpy as np
from sklearn import preprocessing
Category = int | str
_TOL = 1e-6
XRegMode = Literal["timesfm + xreg", "xreg + timesfm"]
def _unnest(nested: Sequence[Sequence[Any]]) -> np.ndarray:
return np.array(list(itertools.chain.from_iterable(nested)))
def _repeat(elements: Iterable[Any], counts: Iterable[int]) -> np.ndarray:
return np.array(
list(
itertools.chain.from_iterable(map(itertools.repeat, elements,
counts))))
def _to_padded_jax_array(x: np.ndarray) -> jax.Array:
if x.ndim == 1:
(i,) = x.shape
di = 2**math.ceil(math.log2(i)) - i
return jnp.pad(x, ((0, di),), mode="constant", constant_values=0.0)
elif x.ndim == 2:
i, j = x.shape
di = 2**math.ceil(math.log2(i)) - i
dj = 2**math.ceil(math.log2(j)) - j
return jnp.pad(x, ((0, di), (0, dj)), mode="constant", constant_values=0.0)
else:
raise ValueError(f"Unsupported array shape: {x.shape}")
class BatchedInContextXRegBase:
"""Helper class for in-context regression covariate formatting.
Attributes:
targets: List of targets (responses) of the in-context regression.
train_lens: List of lengths of each target vector from the context.
test_lens: List of lengths of each forecast horizon.
train_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the context. Their
lengths should match the corresponding lengths in `train_lens`.
train_dynamic_categorical_covariates: Dict of covariate names mapping to the
dynamic categorical covariates of each forecast task on the context. Their
lengths should match the corresponding lengths in `train_lens`.
test_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the horizon. Their
lengths should match the corresponding lengths in `test_lens`.
test_dynamic_categorical_covariates: Dict of covariate names mapping to the
dynamic categorical covariates of each forecast task on the horizon. Their
lengths should match the corresponding lengths in `test_lens`.
static_numerical_covariates: Dict of covariate names mapping to the static
numerical covariates of each forecast task.
static_categorical_covariates: Dict of covariate names mapping to the static
categorical covariates of each forecast task.
"""
def __init__(
self,
targets: Sequence[Sequence[float]],
train_lens: Sequence[int],
test_lens: Sequence[int],
train_dynamic_numerical_covariates: (
Mapping[str, Sequence[Sequence[float]]] | None) = None,
train_dynamic_categorical_covariates: (
Mapping[str, Sequence[Sequence[Category]]] | None) = None,
test_dynamic_numerical_covariates: (
Mapping[str, Sequence[Sequence[float]]] | None) = None,
test_dynamic_categorical_covariates: (
Mapping[str, Sequence[Sequence[Category]]] | None) = None,
static_numerical_covariates: Mapping[str, Sequence[float]] | None = None,
static_categorical_covariates: (Mapping[str, Sequence[Category]] |
None) = None,
) -> None:
"""Initializes with the exogenous covariate inputs.
Here we use model fitting language to refer to the context as 'train' and
the horizon as 'test'. We assume batched inputs. To properly format the
request:
- `train_lens` represents the contexts in the batch. Targets and all train
dynamic covariates should have the same lengths as the corresponding
elements
in `train_lens`. Notice each `train_len` can be different from the exact
length of the corresponding context depending on how much of the context is
used for fitting the in-context model.
- `test_lens` represents the horizon lengths in the batch. All tesdt
dynamic
covariates should have the same lengths as the corresponding elements in
`test_lens`.
- Static covariates should be one for each input.
- For train and test dynamic covariates, they should have the same
covariate
names.
Pass an empty dict {} for a covariate type if it is not present.
Example:
Here is a set of valid inputs whose schema can be used for reference.
```
targets = [
[0.0, 0.1, 0.2],
[0.0, 0.1, 0.2, 0.3],
] # Two inputs in this batch.
train_lens = [3, 4]
test_lens = [2, 5] # Forecast horizons 2 and 5 respectively.
train_dynamic_numerical_covariates = {
"cov_1_dn": [[0.0, 0.5, 1.0], [0.0, 0.5, 1.0, 1.5]],
"cov_2_dn": [[0.0, 1.5, 1.0], [0.0, 1.5, 1.0, 2.5]],
} # Each train dynamic covariate has 3 and 4 elements respectively.
test_dynamic_numerical_covariates = {
"cov_1_dn": [[0.1, 0.6], [0.1, 0.6, 1.1, 1.6, 2.4]],
"cov_2_dn": [[0.1, 1.1], [0.1, 1.6, 1.1, 2.6, 10.0]],
} # Each test dynamic covariate has 2 and 5 elements respectively.
train_dynamic_categorical_covariates = {
"cov_1_dc": [[0, 1, 0], [0, 1, 2, 3]],
"cov_2_dc": [["good", "bad", "good"], ["good", "good", "bad",
"bad"]],
}
test_dynamic_categorical_covariates = {
"cov_1_dc": [[1, 0], [1, 0, 2, 3, 1]],
"cov_2_dc": [["bad", "good"], ["bad", "bad", "bad", "bad", "bad"]],
}
static_numerical_covariates = {
"cov_1_sn": [0.0, 3.0],
"cov_2_sn": [2.0, 1.0],
"cov_3_sn": [1.0, 2.0],
} # Each static covariate has 1 element for each input.
static_categorical_covariates = {
"cov_1_sc": ["apple", "orange"],
"cov_2_sc": [2, 3],
}
```
Args:
targets: List of targets (responses) of the in-context regression.
train_lens: List of lengths of each target vector from the context.
test_lens: List of lengths of each forecast horizon.
train_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the context. Their
lengths should match the corresponding lengths in `train_lens`.
train_dynamic_categorical_covariates: Dict of covariate names mapping to
the dynamic categorical covariates of each forecast task on the context.
Their lengths should match the corresponding lengths in `train_lens`.
test_dynamic_numerical_covariates: Dict of covariate names mapping to the
dynamic numerical covariates of each forecast task on the horizon. Their
lengths should match the corresponding lengths in `test_lens`.
test_dynamic_categorical_covariates: Dict of covariate names mapping to
the dynamic categorical covariates of each forecast task on the horizon.
Their lengths should match the corresponding lengths in `test_lens`.
static_numerical_covariates: Dict of covariate names mapping to the static
numerical covariates of each forecast task.
static_categorical_covariates: Dict of covariate names mapping to the
static categorical covariates of each forecast task.
"""
self.targets = targets
self.train_lens = train_lens
self.test_lens = test_lens
self.train_dynamic_numerical_covariates = (
train_dynamic_numerical_covariates or {})
self.train_dynamic_categorical_covariates = (
train_dynamic_categorical_covariates or {})
self.test_dynamic_numerical_covariates = (test_dynamic_numerical_covariates
or {})
self.test_dynamic_categorical_covariates = (
test_dynamic_categorical_covariates or {})
self.static_numerical_covariates = static_numerical_covariates or {}
self.static_categorical_covariates = static_categorical_covariates or {}
def _assert_covariates(self, assert_covariate_shapes: bool = False) -> None:
"""Verifies the validity of the covariate inputs."""
# Check presence.
if (self.train_dynamic_numerical_covariates and
not self.test_dynamic_numerical_covariates) or (
not self.train_dynamic_numerical_covariates and
self.test_dynamic_numerical_covariates):
raise ValueError(
"train_dynamic_numerical_covariates and"
" test_dynamic_numerical_covariates must be both present or both"
" absent.")
if (self.train_dynamic_categorical_covariates and
not self.test_dynamic_categorical_covariates) or (
not self.train_dynamic_categorical_covariates and
self.test_dynamic_categorical_covariates):
raise ValueError(
"train_dynamic_categorical_covariates and"
" test_dynamic_categorical_covariates must be both present or both"
" absent.")
# Check keys.
for dict_a, dict_b, dict_a_name, dict_b_name in (
(
self.train_dynamic_numerical_covariates,
self.test_dynamic_numerical_covariates,
"train_dynamic_numerical_covariates",
"test_dynamic_numerical_covariates",
),
(
self.train_dynamic_categorical_covariates,
self.test_dynamic_categorical_covariates,
"train_dynamic_categorical_covariates",
"test_dynamic_categorical_covariates",
),
):
if w := set(dict_a.keys()) - set(dict_b.keys()):
raise ValueError(
f"{dict_a_name} has keys not present in {dict_b_name}: {w}")
if w := set(dict_b.keys()) - set(dict_a.keys()):
raise ValueError(
f"{dict_b_name} has keys not present in {dict_a_name}: {w}")
# Check shapes.
if assert_covariate_shapes:
if len(self.targets) != len(self.train_lens):
raise ValueError(
"targets and train_lens must have the same number of elements.")
if len(self.train_lens) != len(self.test_lens):
raise ValueError(
"train_lens and test_lens must have the same number of elements.")
for i, (target, train_len) in enumerate(zip(self.targets,
self.train_lens)):
if len(target) != train_len:
raise ValueError(
f"targets[{i}] has length {len(target)} != expected {train_len}.")
for key, values in self.static_numerical_covariates.items():
if len(values) != len(self.train_lens):
raise ValueError(
f"static_numerical_covariates has key {key} with number of"
f" examples {len(values)} != expected {len(self.train_lens)}.")
for key, values in self.static_categorical_covariates.items():
if len(values) != len(self.train_lens):
raise ValueError(
f"static_categorical_covariates has key {key} with number of"
f" examples {len(values)} != expected {len(self.train_lens)}.")
for lens, dict_cov, dict_cov_name in (
(
self.train_lens,
self.train_dynamic_numerical_covariates,
"train_dynamic_numerical_covariates",
),
(
self.train_lens,
self.train_dynamic_categorical_covariates,
"train_dynamic_categorical_covariates",
),
(
self.test_lens,
self.test_dynamic_numerical_covariates,
"test_dynamic_numerical_covariates",
),
(
self.test_lens,
self.test_dynamic_categorical_covariates,
"test_dynamic_categorical_covariates",
),
):
for key, cov_values in dict_cov.items():
if len(cov_values) != len(lens):
raise ValueError(
f"{dict_cov_name} has key {key} with number of examples"
f" {len(cov_values)} != expected {len(lens)}.")
for i, cov_value in enumerate(cov_values):
if len(cov_value) != lens[i]:
raise ValueError(
f"{dict_cov_name} has key {key} with its {i}-th example"
f" length {len(cov_value)} != expected {lens[i]}.")
def create_covariate_matrix(
self,
one_hot_encoder_drop: str | None = "first",
use_intercept: bool = True,
assert_covariates: bool = False,
assert_covariate_shapes: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Creates target vector and covariate matrices for in context regression.
Here we use model fitting language to refer to the context as 'train' and
the horizon as 'test'.
Args:
one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
use_intercept: Whether to prepare an intercept (all 1) column in the
matrices.
assert_covariates: Whether to assert the validity of the covariate inputs.
assert_covariate_shapes: Whether to assert the shapes of the covariate
inputs when `assert_covariates` is True.
Returns:
A tuple of the target vector, the covariate matrix for the context, and
the covariate matrix for the horizon.
"""
if assert_covariates:
self._assert_covariates(assert_covariate_shapes)
x_train, x_test = [], []
# Numerical features.
for name in sorted(self.train_dynamic_numerical_covariates):
x_train.append(
_unnest(self.train_dynamic_numerical_covariates[name])[:, np.newaxis])
x_test.append(
_unnest(self.test_dynamic_numerical_covariates[name])[:, np.newaxis])
for covs in self.static_numerical_covariates.values():
x_train.append(_repeat(covs, self.train_lens)[:, np.newaxis])
x_test.append(_repeat(covs, self.test_lens)[:, np.newaxis])
if x_train:
x_train = np.concatenate(x_train, axis=1)
x_test = np.concatenate(x_test, axis=1)
# Normalize for robustness.
x_mean = np.mean(x_train, axis=0, keepdims=True)
x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w,
1.0)
x_train = [(x_train - x_mean) / x_std]
x_test = [(x_test - x_mean) / x_std]
# Categorical features. Encode one by one.
one_hot_encoder = preprocessing.OneHotEncoder(
drop=one_hot_encoder_drop,
sparse_output=False,
handle_unknown="ignore",
)
for name in sorted(self.train_dynamic_categorical_covariates.keys()):
ohe_train = _unnest(
self.train_dynamic_categorical_covariates[name])[:, np.newaxis]
ohe_test = _unnest(
self.test_dynamic_categorical_covariates[name])[:, np.newaxis]
x_train.append(np.array(one_hot_encoder.fit_transform(ohe_train)))
x_test.append(np.array(one_hot_encoder.transform(ohe_test)))
for covs in self.static_categorical_covariates.values():
ohe = one_hot_encoder.fit_transform(np.array(covs)[:, np.newaxis])
x_train.append(_repeat(ohe, self.train_lens))
x_test.append(_repeat(ohe, self.test_lens))
x_train = np.concatenate(x_train, axis=1)
x_test = np.concatenate(x_test, axis=1)
if use_intercept:
x_train = np.pad(x_train, ((0, 0), (1, 0)), constant_values=1.0)
x_test = np.pad(x_test, ((0, 0), (1, 0)), constant_values=1.0)
return _unnest(self.targets), x_train, x_test
def fit(self) -> Any:
raise NotImplementedError("Fit is not implemented.")
class BatchedInContextXRegLinear(BatchedInContextXRegBase):
"""Linear in-context regression model."""
def fit(
self,
ridge: float = 0.0,
one_hot_encoder_drop: str | None = "first",
use_intercept: bool = True,
force_on_cpu: bool = False,
max_rows_per_col: int = 0,
max_rows_per_col_sample_seed: int = 42,
debug_info: bool = False,
assert_covariates: bool = False,
assert_covariate_shapes: bool = False,
) -> (list[np.ndarray] | tuple[list[np.ndarray], list[np.ndarray], jax.Array,
jax.Array, jax.Array]):
"""Fits a linear model for in-context regression.
Args:
ridge: A non-negative value for specifying the ridge regression penalty.
If 0 is provided, fallback to ordinary least squares. Note this penalty
is added to the normalized covariate matrix.
one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
use_intercept: Whether to prepare an intercept (all 1) column in the
matrices.
force_on_cpu: Whether to force execution on cpu for accelerator machines.
max_rows_per_col: How many rows to subsample per column. 0 for no
subsampling. This is for speeding up model fitting.
max_rows_per_col_sample_seed: The seed for the subsampling if needed by
`max_rows_per_col`.
debug_info: Whether to return debug info.
assert_covariates: Whether to assert the validity of the covariate inputs.
assert_covariate_shapes: Whether to assert the shapes of the covariate
inputs when `assert_covariates` is True.
Returns:
If `debug_info` is False:
The linear fits on the horizon.
If `debug_info` is True:
A tuple of:
- the linear fits on the horizon,
- the linear fits on the context,
- the flattened target vector,
- the covariate matrix for the context, and
- the covariate matrix for the horizon.
"""
flat_targets, x_train_raw, x_test = self.create_covariate_matrix(
one_hot_encoder_drop=one_hot_encoder_drop,
use_intercept=use_intercept,
assert_covariates=assert_covariates,
assert_covariate_shapes=assert_covariate_shapes,
)
x_train = x_train_raw.copy()
if max_rows_per_col:
nrows, ncols = x_train.shape
if nrows > (w := ncols * max_rows_per_col):
subsample = jax.random.choice(
jax.random.PRNGKey(max_rows_per_col_sample_seed),
nrows,
(w,),
replace=False,
)
x_train = x_train[subsample]
flat_targets = flat_targets[subsample]
device = jax.devices("cpu")[0] if force_on_cpu else None
# Runs jitted version of the solvers which are quicker at the cost of
# running jitting during the first time calling. Re-jitting happens whenever
# new (padded) shapes are encountered.
# Ocassionally it helps with the speed and the accuracy if we force single
# thread execution on cpu for accelerator machines:
# 1. Avoid moving data to accelarator memory.
# 2. Avoid precision loss if any.
with jax.default_device(device):
x_train_raw = _to_padded_jax_array(x_train_raw)
x_train = _to_padded_jax_array(x_train)
flat_targets = _to_padded_jax_array(flat_targets)
x_test = _to_padded_jax_array(x_test)
beta_hat = (jnp.linalg.pinv(
x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]),
hermitian=True,
) @ x_train.T @ flat_targets)
y_hat = x_test @ beta_hat
y_hat_context = x_train_raw @ beta_hat if debug_info else None
outputs = []
outputs_context = []
# Reconstruct the ragged 2-dim batched forecasts from flattened linear fits.
train_index, test_index = 0, 0
for train_index_delta, test_index_delta in zip(self.train_lens,
self.test_lens):
outputs.append(np.array(y_hat[test_index:(test_index +
test_index_delta)]))
if debug_info:
outputs_context.append(
np.array(y_hat_context[train_index:(train_index +
train_index_delta)]))
train_index += train_index_delta
test_index += test_index_delta
if debug_info:
return outputs, outputs_context, flat_targets, x_train, x_test
else:
return outputs
v1/tests/test_timesfm.py
+91
View File
@@ -0,0 +1,91 @@
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from datetime import datetime, timedelta
import numpy as np
import pandas as pd
import pytest
import timesfm
def create_sample_dataframe(
start_date: datetime, end_date: datetime, freq: str = "D"
) -> pd.DataFrame:
"""
Create a sample DataFrame with time series data.
Args:
start_date (datetime): Start date of the time series.
end_date (datetime): End date of the time series.
freq (str): Frequency of the time series (default: "D" for daily).
Returns:
pd.DataFrame: DataFrame with columns 'unique_id', 'ds', and 'ts'.
"""
date_range = pd.date_range(start=start_date, end=end_date, freq=freq)
ts_data = np.random.randn(len(date_range))
df = pd.DataFrame({"unique_id": "ts-1", "ds": date_range, "ts": ts_data})
return df
@pytest.mark.parametrize("context_length", [128, 256, 512])
@pytest.mark.parametrize("prediction_length", [96, 128, 256])
@pytest.mark.parametrize("freq", ["D", "H", "W"])
def test_timesfm_forecast_on_df(
context_length: int,
prediction_length: int,
freq: str,
) -> None:
model = timesfm.TimesFm(
context_len=context_length,
horizon_len=prediction_length,
input_patch_len=32,
output_patch_len=128,
num_layers=20,
model_dims=1280,
backend="cpu",
)
model.load_from_checkpoint(repo_id="google/timesfm-1.0-200m")
end_date = datetime.now()
start_date = end_date - timedelta(days=context_length)
input_df = create_sample_dataframe(start_date, end_date, freq)
forecast_df = model.forecast_on_df(
inputs=input_df,
freq=freq,
value_name="ts",
num_jobs=-1,
)
assert (
len(forecast_df) == prediction_length
), f"Expected forecast length of {prediction_length}, but got {len(forecast_df)}"
assert (
"timesfm" in forecast_df.columns
), "Forecast DataFrame should contain 'timesfm' column"
last_input_date = input_df["ds"].max()
first_forecast_date = forecast_df["ds"].min()
expected_first_forecast_date = last_input_date + pd.Timedelta(1, unit=freq)
assert (
first_forecast_date == expected_first_forecast_date
), f"Forecast should start from {expected_first_forecast_date}, but starts from {first_forecast_date}"
print(
f"Successful forecast with context_length={context_length}, prediction_length={prediction_length}, freq={freq}"
)