tirex.classification

Classification Model interface

class tirex.models.classification.TirexLinearClassifier(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, stratify=True, patience=7, delta=0.001, log_every_n_steps=5, seed=None, class_weights=None, dropout=None)

A PyTorch classifier that combines time series embeddings with a linear classification 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 classification. The embedding backbone is frozen during training, and only the classification head is trained.

Example

>>> import torch
>>> from tirex.models.classification import TirexLinearClassifier
>>>
>>> # Create model with TiRex embeddings
>>> model = TirexLinearClassifier(
...     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.randint(0, 3, (100,))  # 3 classes
>>>
>>> # 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)
>>> probabilities = model.predict_proba(X_test)

• Parameters:
  • data_augmentation (bool)
  • device (str | None)
  • compile (bool)
  • max_epochs (int)
  • lr (float)
  • weight_decay (float)
  • batch_size (int)
  • val_split_ratio (float)
  • stratify (bool)
  • patience (int)
  • delta (float)
  • log_every_n_steps (int)
  • seed (int | None)
  • class_weights (Tensor | None)
  • dropout (float | None)

fit(train_data, val_data=None)

Train the classification head on the provided data.

This method initializes the classification 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 class labels.
  • 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.
• Returns: Dictionary containing final training and validation losses.
• Return type: dict[str, float]

predict(x)

Predict class labels for input time series data.

• Parameters: x (Tensor) – Input tensor of time series data with shape (batch_size, num_variates, seq_len).
• Returns: Predicted class labels with shape (batch_size,).
• Return type: torch.Tensor

predict_proba(x)

Predict class probabilities for input time series data.

• Parameters: x (Tensor) – Input tensor of time series data with shape (batch_size, num_variates, seq_len).
• Returns: Class probabilities with shape (batch_size, num_classes).
• Return type: torch.Tensor

save_model(path)

Save the trained classification head.

This function saves the trained classification 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: TirexLinearClassifier

class tirex.models.classification.TirexRFClassifier(data_augmentation=False, device=None, compile=False, batch_size=512, **rf_kwargs)

A Random Forest classifier that uses time series embeddings as features.

This classifier combines a pre-trained embedding model for feature extraction with a scikit-learn Random Forest classifier. The embedding model generates fixed-size feature vectors from variable-length time series, which are then used to train the Random Forest.

Example

>>> from tirex.models.classification import TirexRFClassifier
>>>
>>> # Create model with custom Random Forest parameters
>>> model = TirexRFClassifier(
...     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.randint(0, 3, (100,))  # 3 classes
>>>
>>> # Train the model
>>> model.fit((X_train, y_train))
>>>
>>> # Make predictions
>>> X_test = torch.randn(20, 1, 128)
>>> predictions = model.predict(X_test)
>>> probabilities = model.predict_proba(X_test)

• Parameters:
  • data_augmentation (bool)
  • device (str | None)
  • compile (bool)
  • batch_size (int)

fit(train_data)

Train the Random Forest classifier 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 class labels.
• Return type: None

save_model(path)

This method saves the trained Random Forest classifier 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 classifier 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, ready for inference.
• Return type: TirexRFClassifier

predict(x)

Predict class labels 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 class labels with shape (batch_size,).
• Return type: torch.Tensor

predict_proba(x)

Predict class probabilities for input time series data.

• Parameters: x (Tensor) – Input time series data as torch.Tensor with shape (batch_size, num_variates, seq_len).
• Returns: Class probabilities with shape (batch_size, num_classes).
• Return type: torch.Tensor

class tirex.models.classification.TirexGBMClassifier(data_augmentation=False, device=None, compile=False, batch_size=512, early_stopping_rounds=10, min_delta=0.0, val_split_ratio=0.2, stratify=True, **lgbm_kwargs)

A Gradient Boosting classifier that uses time series embeddings as features.

This classifier combines a pre-trained embedding model for feature extraction with a Gradient Boosting classifier.

Example

>>> from tirex.models.classification import TirexGBMClassifier
>>>
>>> # Create model with custom LightGBM parameters
>>> model = TirexGBMClassifier(
...     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.randint(0, 3, (100,))  # 3 classes
>>>
>>> # Train the model
>>> model.fit((X_train, y_train))
>>>
>>> # Make predictions
>>> X_test = torch.randn(20, 1, 128)
>>> predictions = model.predict(X_test)
>>> probabilities = model.predict_proba(X_test)

• Parameters:
  • data_augmentation (bool)
  • device (str | None)
  • compile (bool)
  • batch_size (int)
  • early_stopping_rounds (int | None)
  • min_delta (float)
  • val_split_ratio (float)
  • stratify (bool)

fit(train_data, val_data=None)

Train the LightGBM classifier on embedded time series data.

This method generates embeddings for the training data using the embedding model, then trains the LightGBM classifier 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 class labels.
  • val_data (tuple[Tensor, Tensor] | None) – Optional tuple of (X_val, y_val) for validation where X_train is the input time series data (torch.Tensor) and y_train is a torch.Tensor of class labels. If None, validation data will be automatically split from train_data (20% split).
• Return type: None

save_model(path)

This method saves the trained LightGBM classifier 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 classifier 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, ready for inference.
• Return type: TirexGBMClassifier

predict(x)

Predict class labels 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 class labels with shape (batch_size,).
• Return type: torch.Tensor

predict_proba(x)

Predict class probabilities for input time series data.

• Parameters: x (Tensor) – Input time series data as torch.Tensor with shape (batch_size, num_variates, seq_len).
• Returns: Class probabilities with shape (batch_size, num_classes).
• Return type: torch.Tensor