ddrg/ChronosBERT-Default

ChronosBERT

ChronosBERT is a pretrained BERT-style bidirectional encoder-only model for time series representation learning.

Model Description

• Architecture: BERT-Small (encoder-only)
• Tokenizer: Chronos MeanScaleUniformBins
• Pretraining objective: Masked Language Modeling
• Number of parameters: Approximately 29M

ChronosBERT combines the Chronos tokenizer with a BERT architecture and a masked language modeling (MLM) pretraining objective.

Time series are normalized and discretized into tokens using the Chronos tokenizer. These tokens are then processed by a bidirectional BERT encoder, which is trained to reconstruct randomly masked tokens. This allows the model to learn general time series representations that can be transferred to downstream tasks.

ChronosBERT can be fine-tuned for individual downstream tasks.

GitHub

• Repository: https://github.com/JuliusHanusch/ChronosBERT

Fine-tuning for Time Series Classification

The following example demonstrates how to fine-tune ChronosBERT on a time series classification dataset.

Setup

Clone the ChronosBERT repository:

git clone https://github.com/JuliusHanusch/ChronosBERT.git

For details on the environment setup and installation, refer to the GitHub repository.

Fine-tuning for Time Series Classification

The following example demonstrates how to fine-tune ChronosBERT on the ArrowHead dataset from the UCR Time Series Classification Archive.

Prepare the Dataset

Download the UCR Time Series Classification Archive, extract it, and locate the ArrowHead dataset.

The repository also provides preprocessing scripts for extracting the UCR classification datasets:

python3 data/finetuning/UCR_extracted/UCR_classification_data.py

Load the Dataset

import numpy as np
import pandas as pd

# Load training and test data
train = pd.read_csv("ArrowHead_TRAIN.tsv", sep="\t", header=None).values
test = pd.read_csv("ArrowHead_TEST.tsv", sep="\t", header=None).values

X_train = train[:, 1:].astype(np.float32)
y_train = train[:, 0].astype(int)

X_test = test[:, 1:].astype(np.float32)
y_test = test[:, 0].astype(int)

# Map class labels to consecutive integers
classes = np.unique(y_train)
label_map = {v: i for i, v in enumerate(classes)}

y_train = np.vectorize(label_map.get)(y_train)
y_test = np.vectorize(label_map.get)(y_test)

num_labels = len(classes)

Load ChronosBERT

import sys
sys.path.append("ChronosBERT/chronos_pkg/src")

from chronos_pkg.src.chronos import ChronosPipeline

pipeline = ChronosPipeline.from_pretrained(
    "ddrg/ChronosBERT-Default",
    task="classification",
    num_labels=num_labels,
    dropout_head=0.1,
    TrainInnerModel=True,
)

model = pipeline.model
tokenizer = pipeline.tokenizer

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

Prepare the Dataset

class ClassificationDataset(Dataset):

    def __init__(self, X, y, tokenizer):
        self.X = X
        self.y = y
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        series = torch.tensor(
            self.X[idx],
            dtype=torch.float32
        )

        token_ids, attention_mask, _ = (
            self.tokenizer.context_input_transform(
                series.unsqueeze(0)
            )
        )

        return {
            "input_ids": token_ids.squeeze(0),
            "attention_mask": attention_mask.squeeze(0),
            "labels": torch.tensor(
                self.y[idx],
                dtype=torch.long
            ),
        }


train_dataset = ClassificationDataset(
    X_train, y_train, tokenizer
)

test_dataset = ClassificationDataset(
    X_test, y_test, tokenizer
)

train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True
)

test_loader = DataLoader(
    test_dataset,
    batch_size=32
)

Finetune the Model

from torch.optim import AdamW

optimizer = AdamW(
    model.parameters(),
    lr=1e-4
)

num_epochs = 10

for epoch in range(num_epochs):

    model.train()
    total_loss = 0

    for batch in train_loader:

        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )

        loss = outputs["loss"]

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

Evaluate the Model

model.eval()

correct = 0
total = 0

with torch.no_grad():

    for batch in test_loader:

        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

        predictions = outputs["logits"].argmax(dim=-1)

        correct += (predictions == labels).sum().item()
        total += labels.size(0)

accuracy = correct / total

print(f"Accuracy: {accuracy:.4f}")