iTransformer Documentation

Install the iTransformer package using pip.

$ pip install iTransformer

The iTransformer class implements the SOTA time series forecasting architecture. It accepts a time series tensor and returns a dictionary of predictions for various lengths.

Input Shape: (batch, lookback_len, num_variates) Output Format: Dict[int, Tensor[batch, pred_length, variate]] where the key is the prediction length.

import torch
from iTransformer import iTransformer

model = iTransformer(
    num_variates = 137,
    lookback_len = 96,
    dim = 256,
    depth = 6,
    heads = 8,
    dim_head = 64,
    pred_length = (12, 24, 36, 48), # Can be a single int or a tuple of lengths
    num_tokens_per_variate = 1,     # Experimental: projects each variate to multiple tokens for granular time attention
    use_reversible_instance_norm = True
)

time_series = torch.randn(2, 96, 137)  # (batch, lookback len, variates)
preds = model(time_series)

The iTransformerFFT model is an experimental variant that includes Fourier tokens. The FFT of the time series is projected into its own tokens and attended alongside the variate tokens, then spliced back in at the end.

import torch
from iTransformer import iTransformerFFT

model = iTransformerFFT(
    num_variates = 137,
    lookback_len = 96,
    dim = 256,
    depth = 6,
    heads = 8,
    dim_head = 64,
    pred_length = (12, 24, 36, 48),
    num_tokens_per_variate = 1,
    use_reversible_instance_norm = True
)

time_series = torch.randn(2, 96, 137)
preds = model(time_series)

The iTransformer2D model is an improvised version that performs granular attention across both variates and time tokens. This is achieved by specifying num_time_tokens, which determines the patch size (calculated as lookback_len // num_time_tokens).

import torch
from iTransformer import iTransformer2D

model = iTransformer2D(
    num_variates = 137,
    num_time_tokens = 16,               # Number of time tokens
    lookback_len = 96,
    dim = 256,
    depth = 6,
    heads = 8,
    dim_head = 64,
    pred_length = (12, 24, 36, 48),
    use_reversible_instance_norm = True
)

time_series = torch.randn(2, 96, 137)
preds = model(time_series)