Using the IBM Granite models for time series forecasting

Forecasting in time series analysis allows data scientists to identify patterns by using machine learning and then generate forecasts about the future. Deep learning for forecasting is beginning to show promise when compared to benchmarks from more traditional statistical methods such as ARIMA.

In this tutorial, we'll use the IBM Granite TinyTimeMixer (TTM) model and the IBM Time Series Foundation Model libraries in Python to test how TTM does in two different forecasting tasks. Zero-shot learning is when the model is not trained on the data it’s trying to predict. It’s an interesting test of how well our model can detect and respond to the patterns present in the time series. After that, we’ll fine tune the model to see whether there’s a performance boost.

More about time series foundation models

Foundation models for time series data are similar to other forms of generative AI models that are trained on large-scale time series data sets and can output either deterministic or probabilistic forecasts. Unfortunately, though, large language models (LLMs) don't perform especially well in time series forecasting tasks (see this paper, Are Language Models Actually Useful for Time Series Forecasting?). But LLMs aren't the only types of foundation models. Foundation models have been built specifically for time series forecasts, such as Moirai, TimeGPT-1 and TimesFM.

The Granite TinyTimeMixer (TTM) is another compact time-series foundation model for forecasting, open-sourced by IBM Research. With fewer than 1 million parameters, TTM introduces the notion of “tiny” pretrained models for time-series forecasting. Unlike LLM-based forecasting models, TTM is a fast and compact model (starting from 1 million parameters) with effective transfer learning capabilities, trained exclusively on public TS datasets. TTM, based on the lightweight TSMixer architecture, incorporates innovations such as adaptive patching, diverse resolution sampling and resolution prefix tuning to handle pretraining on varied dataset resolutions with minimal model capacity.

TinyTimeMixer doesn't use self-attention heads like a transformer model does. Instead, it uses what's called gated attention to control simpler blocks of perceptrons that correlate multiple variables in a time series. The goal of this architecture is to keep the model from overcomplicating what it pays attention to and to make training and fine-tuning less computationally expensive.

The version of TTM we'll use in this tutorial is pretrained using a 1 billion parameter pretraining dataset that can be pre-trained in only 24–30 hours. This scenario is notably faster than transformer-based approaches that often take days to weeks to train.

Prerequisites

You need an IBM Cloud account to create a watsonx.ai project.

Steps

Step 1. Set up your environment

In this step, we’ll guide you through creating an IBM account to access Jupyter Notebooks.

1. Log in to watsonx.ai using your IBM Cloud account.

2. Create a watsonx.ai project.

3. Create a Jupyter Notebook.

   Choose "Runtime 23.1 on Python (2 vCPU 8 GB RAM)" to define the configuration.

Step 2. Clone the repository and install libraries

You can download the notebook from Github.

To use the TTM model, we'll clone the open source IBM Time Series Foundation Model (TSFM) library from the Github repository and install it on the watsonx.ai instance.

# Clone the ibm and tsfm
! git clone --depth 1 --branch v0.2.9 https://github.com/ibm-granite/granite-tsfm
# Change directory. Move inside the tsfm repo.
%cd tsfm
# Install the tsfm library and scikit-learn
! pip install ".[notebooks]" scikit-learn
%cd ../

Step 3. Import the libraries

TSFM relies on several Hugging Face libraries that we'll use to load our models, fine-tune them and get forecasts.

import os
import math
import tempfile
import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import OneCycleLR
from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed
import numpy as np
import pandas as pd

# TSFM libraries
from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction
from tsfm_public.toolkit.callbacks import TrackingCallback

Step 4. Load the data

We'll use the Beijing Air Quality dataset. This hourly data set contains the multivariate time series data of the PM2.5 data of the US Embassy in Beijing. It also includes meteorological data from the Beijing Capital International Airport that we can use as exogenous variables for our forecasts. Our goal is to forecast PM2.5, stored in our target_column, using past values of the data and exogenous variables, stored in the observable_columns.

pollution_data = pd.read_csv(
    "pollution.csv"
)

timestamp_column = "date"
target_columns = ["pm2.5"]
observable_columns = ["DEWP","TEMP","PRES", "Iws"]

pollution_data['date'] = pd.to_datetime(pollution_data[['year','month','day', 'hour']])

Because there are sensor errors leading to NaN (not a number) values in the time series, we'll interpolate the values to fill in any missing values.

#deal with NA
pollution_data['pm2.5'] = pollution_data['pm2.5'].interpolate()

When we look at the PM2.5 and windspeed, we can see that it's going to be a difficult time series dataset to work with because of the significant changes between data points and time points.

pollution_data[:1000].plot(x="date", y=["pm2.5", "Iws"], figsize=(20,5))

This returns the two time series plotted together:

Step 5: Prepare the data

Next, we prepare our data to do zero-shot prediction and then fine-tuning.

# Set seed for reproducibility

SEED = 42
set_seed(SEED)

# Results dir
OUT_DIR = "ttm_finetuned_models/"

# Forecasting parameters
context_length = 512 # TTM can use 512 time points into the past
forecast_length = 96 # TTM can predict 96 time points into the future

Next we configure the train, evaluation, and testing splits in the data set. We'll use 80% of the data for training, 10% for evaluation, and a further 10% for testing. This data is stored in the dictionary object that we use later.

from tsfm_public import (
    TimeSeriesPreprocessor,
    TinyTimeMixerForPrediction,
    TrackingCallback,
    count_parameters,
    get_datasets,
)

data_length = len(pollution_data)

train_start_index = 0
train_end_index = round(data_length * 0.8)

# we shift the start of the evaluation period back by context length so that
# the first evaluation timestamp is immediately following the training data
eval_start_index = round(data_length * 0.8)
eval_end_index = round(data_length * 0.9)

test_start_index = round(data_length * 0.9)
test_end_index = data_length

split_config = {
                "train": [0, train_end_index],
                "valid": [eval_start_index, eval_end_index],
                "test": [test_start_index,test_end_index],
            }

Finally, we specify which columns we'll use and configure a TimeSeriesPreprocessor instance to scale our values to a normalized scale. Once we've configured the TimeSeriesPreprocessor, we use it to generate train, validation, and test data sets with the get_datasets method. This method uses the split configuration to create torch vectors from our data set that can be used for training and evaluation. It also configures the TimeSeriesPreprocessor to properly scale the data down to normalized values that are easier for our model to interpret and predict.

column_specifiers = {
    "timestamp_column": timestamp_column,
    "target_columns": target_columns,
    "observable_columns": observable_columns
}

tsp = TimeSeriesPreprocessor(
    **column_specifiers,
    context_length=context_length,
    prediction_length=forecast_length,
    scaling=True,
    encode_categorical=True,
    scaler_type="standard",
)

# this gets torch vectors for training. For test eval we need a Pandas DF
train_dataset, valid_dataset, test_dataset = get_datasets(
    tsp, pollution_data, split_config
)

Now, we have torch data sets ready to use in testing and training.

Step 6. Load the TTM model and evaluate it

Let's see how the pretrained TTM model does with zero-shot forecasting. We load the model by using TinyTimeMixerForPrediction.from_pretrained() and then configure a HuggingFace trainer to evaluate the test portion of the pollution dataset by calling evaluate()

zeroshot_model = TinyTimeMixerForPrediction.from_pretrained("ibm/TTM", revision="main", prediction_filter_length=24)

# zeroshot_trainer
zeroshot_trainer = Trainer(
    model=zeroshot_model,
)

zeroshot_trainer.evaluate(test_dataset)  #note that this is the Torch dataset created by get_datasets(), not a Pandas DataFrame

This example returns the mean absolute error (MAE):

{'eval_loss': 0.29865384101867676,
 'eval_model_preparation_time': 0.019,
 'eval_runtime': 50.259,
 'eval_samples_per_second': 95.485,
 'eval_steps_per_second': 11.938}

The result shows us that the eval_loss is approximately 0.298. We'll compare this result to the fine-tuned value in a later step.

Step 7. Use the TTM model for forecasting

To use our TTM model for forecasting, we need to create a TimeSeriesForecastingPipeline object that will pass data to our model and generate predictions. We then have our TimeSeriesPreprocessor preprocess the data before we pass it to the TimeSeriesForecastingPipeline.

If you were building a production system that read in real-world data and created predictions from it, you would create a TimeSeriesForecastingPipeline object and pass new data to that to generate forecasts. In this tutorial, for testing purposes, we're just passing the entire test data set to generate predictions for each timestamp in the data set.

from tsfm_public.toolkit.util import select_by_index
from tsfm_public.toolkit.time_series_forecasting_pipeline import TimeSeriesForecastingPipeline
from tsfm_public.toolkit.visualization import plot_ts_forecasting

zs_forecast_pipeline = TimeSeriesForecastingPipeline(
    model=zeroshot_model,
    device="cpu",
    timestamp_column=timestamp_column,
    id_columns=[],
    target_columns=target_columns,
    freq="1h"
)

Now it's time to pass our test data and get results. Note that the TimeSeriesForecastingPipeline takes a Pandas DataFrame. This method is different from the evaluate() method of the trainer, which takes a torch array.

zs_forecast = zs_forecast_pipeline(tsp.preprocess(pollution_data[test_start_index:test_end_index]))

When we look at the returned forecast objects, it's important to note that there are 24 of them in the pm2.5_predictions field. This forecast is because we asked the model to return a window with 24 timestamps worth of predictions for each prediction. To compare the predictions to the actual values, then we can simply plot the predictions and the actual values:

spot_check_index = 200
fcast_df = pd.DataFrame({"pred":zs_forecast.loc[spot_check_index]['pm2.5_prediction'], "actual":zs_forecast.loc[spot_check_index]['pm2.5'][:24]})
fcast_df.plot()

This example generates the following plot:

We can see that the zero-shot model doesn't track the drop in PM2.5 values which is reasonable because it's driven by external factors.

Step 8. Evaluating the forecast results

To evaluate our forecast results, we 'll see how they do against the actual PM2.5 readings 12 hours later.

def compare_forecast(forecast, date_col, prediction_col, actual_col, hours_out):
  comparisons = pd.DataFrame()
  comparisons[date_col] = forecast[date_col]
  actual = []
  pred = []

  for i in range(len(forecast)):
    pred.append(forecast[prediction_col].values[i][hours_out - 1]) # prediction for next day
    actual.append(forecast[actual_col].values[i][hours_out - 1])

  comparisons['actual'] = actual
  comparisons['pred'] = pred

  return comparisons

Now we can use this method to compare our forecasted PM2.5 values to the actual PM2.5 values. We pass the forecast and actual data frames, the fields that we'll need to compare, and the time offset for our predictions. We might use mean absolute error (MAE) or mean absolute percentage error (MAPE) but in our case, with so many very small or negative values the root mean squared error (RMSE) will work well.

half_day_out_predictions = compare_forecast(zs_forecast, "date", "pm2.5_prediction", "pm2.5", 12)

out = half_day_out_predictions[[not np.isnan(x).any() for x in half_day_out_predictions['actual']]]

rms = '{:.10f}'.format(mean_squared_error(half_day_out_predictions['actual'], half_day_out_predictions['pred'], squared=False))
half_day_out_predictions.plot(x="date", y=["pred", "actual"], figsize=(20,5), title=str(rms))

This example shows an RMSE of 0.844 and gives us the following plot:

We can see how the zero-shot model doesn't predict the spikes in PM2.5 and overcorrects to account for them.

Step 9. Fine tuning the TTM model

Now let's fine-tune the model weights of a version of the TTM model with our training data and then evaluate it on the test set. First, we need to configure the TrainingArguments. This is a long function call but boils down to simply configuring the learning rate, the number of epochs to train for, the batch size and the early stopping.

# Important parameters
learning_rate = 0.001
num_epochs = 20
batch_size = 32

finetune_forecast_args = TrainingArguments(
    output_dir=os.path.join(OUT_DIR, "output"),
    overwrite_output_dir=True,
    learning_rate=learning_rate,
    num_train_epochs=num_epochs,
    do_eval=True,
    evaluation_strategy="epoch",
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    dataloader_num_workers=8,
    save_strategy="epoch",
    logging_strategy="epoch",
    save_total_limit=1,
    logging_dir=os.path.join(OUT_DIR, "logs"),  # Specify a logging directory
    load_best_model_at_end=True,  # Load the best model when training ends
    metric_for_best_model="eval_loss",  # Metric to monitor for early stopping
    greater_is_better=False,  # For loss
)

Because our model might overfit, we create an early stopping function. This is also where we configure the scheduler and optimizer that the Trainer will use.

# Create the early stopping callback
early_stopping_callback = EarlyStoppingCallback(
    early_stopping_patience=2,  # Number of epochs with no improvement after which to stop
    early_stopping_threshold=0.001,  # Minimum improvement required to consider as improvement
)
tracking_callback = TrackingCallback()

# Optimizer and scheduler
optimizer = AdamW(finetune_forecast_model.parameters(), lr=learning_rate)
scheduler = OneCycleLR(
    optimizer,
    learning_rate,
    epochs=num_epochs,
    steps_per_epoch=math.ceil(len(train_dataset) / (batch_size)),
)

Once we have all our training parameters configured, we can begin fine-tuning our model. We pass all the configuration parameters to the trainer along with the model that we want to fine-tune and the train and evaluation data sets that we configured earlier. When possible, we want to use early stopping to keep our training process from overfitting our training data. The call to train() starts the process off. You'll notice that the TTM model fine-tunes very quickly, in about 15 minutes just on the CPU and much faster if you use a GPU.

finetune_forecast_trainer = Trainer(
    model=finetune_forecast_model,
    args=finetune_forecast_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    callbacks=[early_stopping_callback, tracking_callback],
    optimizers=(optimizer, scheduler),
)

# Fine tune
finetune_forecast_trainer.train()

Now we have a fine-tuned model that we can evaluate.

Step 10. Evaluating the fine-tuned TTM model

We'll evaluate the fine-tuned TTM model in the same way as the zero-shot version of the model. Our first step will be to use the evaluate() method of the trainer instance:

finetune_forecast_trainer.evaluate(test_dataset) #note that this is the torch dataset, not a Pandas DataFrame

This example will return:

{'eval_loss': 0.24936321377754211,
 'eval_runtime': 1.8753,
 'eval_samples_per_second': 2559.1,
 'eval_steps_per_second': 20.264,
 'epoch': 3.0}

We can see that the eval_loss is 0.249 or 0.05 better than the model that was not fine-tuned. This is a small improvement but a significant one.

Now create a ForecastingPipeline to use the fine-tuned model and the exogenous variables by passing the observable_columns to the pipeline.

forecast_pipeline = TimeSeriesForecastingPipeline(
    model=finetune_forecast_model,
    device="cpu",
    timestamp_column=timestamp_column,
    id_columns=[],
    target_columns=target_columns,
    observable_columns=observable_columns,
    freq="1h"
)

forecasts = forecast_pipeline(test_data)

fcast_df = pd.DataFrame({"pred":forecasts.loc[50]['pm2.5_prediction'], "actual":forecasts.loc[50]['pm2.5'][:24]})
fcast_df.plot()

A spot check on the forecasting from the finetuned model shows that it does a better job of capturing the downward trend of the PM2.5 data:

Let's see how our model does on forecasting 12 hours out:

half_day_out_predictions = compare_forecast(forecasts, "date", "pm2.5_prediction", "pm2.5", 12)

out = half_day_out_predictions[[not np.isnan(x).any() for x in half_day_out_predictions['actual']]]

rms = '{:.10f}'.format(mean_squared_error(out['actual'], out['pred'], squared=False))
half_day_out_predictions.plot(x="date", y=["pred", "actual"], figsize=(20,5), title="RMSE:" + str(rms))

This gives us the following graph:

The RMSE here is 0.80, we can see that our fine-tuning has given us some improvement on a very difficult data set. The forecast corrects much more quickly than the zero-shot version because it has access to the training data and the covariate exogenous variables.

Summary and next steps

In this tutorial you learned about using foundation models for forecasting, an area of research in artificial intelligence with important applications in healthcare, predictive analytics, anomaly detection, and many other fields and tasks.

You used the IBM Granite TinyTimeMixer model, a foundation model trained on time series data specifically built for forecasting. You used the base model to perform zero-shot forecasting on an air pollution data set. Then, you fine-tuned the TTM model data with training data from the air pollution data set. You then generated forecasts from that fine-tuned model and compared its performance to the zero-shot data set.

You can learn more about the Granite Time Series Models at their Hugging Face site.

Try watsonx for free

Build an artificial intelligence (AI) strategy using AI for Forecasting in your business on one collaborative AI and data platform called IBM watsonx, which combines new generative AI capabilities powered by foundation models and traditional machine learning into a powerful platform spanning the AI lifecycle. With watsonx.ai, you can train, validate, tune and deploy models with ease and build AI applications in a fraction of the time with a fraction of the data.

Try watsonx.ai, the next-generation studio for AI builders.

Next steps

Explore more articles and tutorials about watsonx on IBM Developer.