tirex.classification
Classification Model interface
class tirex.models.classification.TirexClassifierTorch(data_augmentation=False, device=None, max_epochs=50, 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 TirexClassifierTorch
>>>
>>> # Create model with TIREX embeddings
>>> model = TirexClassifierTorch(
... 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:
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.
- 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: TirexClassifierTorch
class tirex.models.classification.TirexRFClassifier(data_augmentation=False, device=None, 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
>>> import numpy as np
>>> 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)
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 numpy array of class labels. - Return type:
None
predict(x)
Predict class labels for input time series data.
- Parameters:
x (
Tensor) – Input time series data as torch.Tensor or np.ndarray 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 or np.ndarray with shape (batch_size, num_variates, seq_len). - Returns: Class probabilities with shape (batch_size, num_classes).
- Return type: torch.Tensor
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