tirex.regression
Regression Model interface
class tirex.models.regression.TirexLinearRegressor(data_augmentation=False, device=None, compile=False, max_epochs=10, lr=1e-4, weight_decay=0.01, batch_size=512, val_split_ratio=0.2, patience=7, delta=0.001, log_every_n_steps=5, seed=None, dropout=None)
A PyTorch regressor that combines time series embeddings with a linear regression head.
This model uses a pre-trained TiRex embedding model to generate feature representations from time series data, followed by a linear layer (with optional dropout) for regression. The embedding backbone is frozen during training, and only the regression head is trained.
Example
>>> import torch
>>> from tirex.models.regression import TirexLinearRegressor
>>>
>>> # Create model with TiRex embeddings
>>> model = TirexLinearRegressor(
... data_augmentation=True,
... max_epochs=2,
... lr=1e-4,
... batch_size=32
... )
>>>
>>> # Prepare data
>>> X_train = torch.randn(100, 1, 128) # 100 samples, 1 number of variates, 128 sequence length
>>> y_train = torch.randn(100, 1) # target values
>>>
>>> # Train the model
>>> metrics = model.fit((X_train, y_train))
Epoch 1, Train Loss: ...
>>> # Make predictions
>>> X_test = torch.randn(20, 1, 128)
>>> predictions = model.predict(X_test)
- Parameters:
fit(train_data, val_data=None)
Train the regression head on the provided data.
This method initializes the regression head based on the data dimensions, then trains it on provided data. The embedding model remains frozen.
- Parameters:
- train_data (
tuple[Tensor,Tensor]) – Tuple of (X_train, y_train) where X_train is the input time series data and y_train are the corresponding target values. - val_data (
tuple[Tensor,Tensor] |None) – Optional tuple of (X_val, y_val) for validation. If None and val_split_ratio > 0, validation data will be split from train_data.
- train_data (
- Returns: Dictionary containing final training and validation losses.
- Return type: dict[str, float]
predict(x)
Predict values for input time series data.
- Parameters:
x (
Tensor) – Input tensor of time series data with shape (batch_size, num_variates, seq_len). - Returns: Predicted values with shape (batch_size, 1).
- Return type: torch.Tensor
save_model(path)
Save the trained regression head.
This function saves the trained regression head weights (.pt format), embedding configuration, model dimensions, and device information. The embedding model itself is not saved as it uses a pre-trained backbone that can be reloaded.
- Parameters:
path (
str) – File path where the model should be saved (e.g., ‘model.pt’). - Return type:
None
classmethod load_model(path)
Load a saved model from file.
This reconstructs the model architecture and loads the trained weights from a checkpoint file created by save_model().
- Parameters:
path (
str) – File path to the saved model checkpoint. - Returns: The loaded model with trained weights, ready for inference.
- Return type: TirexLinearRegressor
class tirex.models.regression.TirexRFRegressor(data_augmentation=False, device=None, compile=False, batch_size=512, **rf_kwargs)
A Random Forest regressor that uses time series embeddings as features.
This regressor combines a pre-trained embedding model for feature extraction with a scikit-learn Random Forest regressor. The embedding model generates fixed-size feature vectors from variable-length time series, which are then used to train the Random Forest.
Example
>>> import torch
>>> from tirex.models.regression import TirexRFRegressor
>>>
>>> # Create model with custom Random Forest parameters
>>> model = TirexRFRegressor(
... data_augmentation=True,
... n_estimators=50,
... max_depth=10,
... random_state=42
... )
>>>
>>> # Prepare data (can use NumPy arrays or PyTorch tensors)
>>> X_train = torch.randn(100, 1, 128) # 100 samples, 1 number of variates, 128 sequence length
>>> y_train = torch.randn(100,) # target values
>>>
>>> # Train the model
>>> model.fit((X_train, y_train))
>>>
>>> # Make predictions
>>> X_test = torch.randn(20, 1, 128)
>>> predictions = model.predict(X_test)
fit(train_data)
Train the Random Forest regressor on embedded time series data.
This method generates embeddings for the training data using the embedding model, then trains the Random Forest on these embeddings.
- Parameters:
train_data (
tuple[Tensor,Tensor]) – Tuple of (X_train, y_train) where X_train is the input time series data (torch.Tensor) and y_train is a torch.Tensor of target values. - Return type:
None
save_model(path)
This method saves the trained Random Forest regressor head and embedding information in joblib format
- Parameters:
path (
str) – File path where the model should be saved (e.g., ‘model.joblib’). - Return type:
None
classmethod load_model(path)
Load a saved model from file.
This reconstructs the model with the embedding configuration and loads the trained Random Forest regressor from a checkpoint file created by save_model().
- Parameters:
path (
str) – File path to the saved model checkpoint. - Returns: The loaded model with trained Random Forest regressor, ready for inference.
- Return type: TirexRFRegressor
predict(x)
Predict values for input time series data.
- Parameters:
x (
Tensor) – Input time series data as torch.Tensor with shape (batch_size, num_variates, seq_len). - Returns: Predicted values with shape (batch_size,).
- Return type: torch.Tensor
class tirex.models.regression.TirexGBMRegressor(data_augmentation=False, device=None, compile=False, batch_size=512, early_stopping_rounds=10, min_delta=0.0, val_split_ratio=0.2, **lgbm_kwargs)
A Gradient Boosting regressor that uses time series embeddings as features.
This regressor combines a pre-trained embedding model for feature extraction with a Gradient Boosting regressor.
Example
>>> import torch
>>> from tirex.models.regression import TirexGBMRegressor
>>>
>>> # Create model with custom LightGBM parameters
>>> model = TirexGBMRegressor(
... data_augmentation=True,
... n_estimators=50,
... random_state=42
... )
>>>
>>> # Prepare data (can use NumPy arrays or PyTorch tensors)
>>> X_train = torch.randn(100, 1, 128) # 100 samples, 1 number of variates, 128 sequence length
>>> y_train = torch.randn(100,) # target values
>>>
>>> # Train the model
>>> model.fit((X_train, y_train))
>>>
>>> # Make predictions
>>> X_test = torch.randn(20, 1, 128)
>>> predictions = model.predict(X_test)
- Parameters:
fit(train_data, val_data=None)
Train the LightGBM regressor on embedded time series data.
This method generates embeddings for the training data using the embedding model, then trains the LightGBM regressor on these embeddings.
- Parameters:
- train_data (
tuple[Tensor,Tensor]) – Tuple of (X_train, y_train) where X_train is the input time series data (torch.Tensor) and y_train is a torch.Tensor of target values. - val_data (
tuple[Tensor,Tensor] |None) – Optional tuple of (X_val, y_val) for validation where X_val is the input time series data (torch.Tensor) and y_val is a torch.Tensor of target values. If None, validation data will be automatically split from train_data (20% split).
- train_data (
- Return type:
None
save_model(path)
This method saves the trained LightGBM regressor head (joblib format) and embedding information.
- Parameters:
path (
str) – File path where the model should be saved (e.g., ‘model.joblib’). - Return type:
None
classmethod load_model(path)
Load a saved model from file.
This reconstructs the model with the embedding configuration and loads the trained LightGBM regressor from a checkpoint file created by save_model().
- Parameters:
path (
str) – File path to the saved model checkpoint. - Returns: The loaded model with trained Gradient Boosting regressor, ready for inference.
- Return type: TirexGBMRegressor
predict(x)
Predict values for input time series data.
- Parameters:
x (
Tensor) – Input time series data as torch.Tensor with shape (batch_size, num_variates, seq_len). - Returns: Predicted values with shape (batch_size,).
- Return type: torch.Tensor