ddrg/ChronosBERT-Bolt-Default

ChronosBERT-Bolt

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

Model Description

• Architecture: BERT-Small (encoder-only)
• Input representation: Chronos-Bolt patch-based representation
• Pretraining objective: Masked Language Modeling
• Number of parameters: Approximately 29M

ChronosBERT-Bolt combines the patch-based time series representation of Chronos-Bolt with a BERT architecture and a masked language modeling (MLM) pretraining objective.

Time series are divided into patches using the Chronos-Bolt representation. These patches are then processed by a bidirectional BERT encoder, which is trained to reconstruct randomly masked patches of the input. This allows the model to learn general time series representations that can be transferred to downstream tasks.

ChronosBERT-Bolt 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-Bolt on a time series classification dataset.

Load the Dataset

Download the UCR Time Series Classification Archive and extract the ArrowHead 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.array([label_map[y] for y in y_train])
y_test = np.array([label_map[y] for y in y_test])

num_labels = len(classes)

Load ChronosBERT-Bolt

ChronosBERT-Bolt uses the pretrained Chronos-Bolt forecasting model as its backbone. A classification model is constructed using the pretrained configuration, and the pretrained weights are transferred to it. The classification head is initialized separately.

import torch
from chronos_pkg.src.chronos.chronos_bolt import (
    ChronosBoltModelForForecasting,
)
from chronos_pkg.src.chronos.chronos_bolt_classification import (
    ChronosBoltModelForClassification,
)

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

# Load pretrained model
pretrained_model = ChronosBoltModelForForecasting.from_pretrained(
    "ddrg/ChronosBERT-Bolt-Default"
)

# Disable random patch augmentation
config = pretrained_model.chronos_config

for attribute in [
    "masking_prob",
    "patch_masking_prob",
    "perturbation_prob",
    "patch_perturbation_prob",
]:
    if hasattr(config, attribute):
        setattr(config, attribute, 0.0)

# Get context length
context_length = config.context_length

config = pretrained_model.config


# Initialize classification model
model = ChronosBoltModelForClassification(
    config=config,
    num_labels=num_labels,
    dropout_head=0.1,
    pooling="mean",
)

# Transfer pretrained backbone weights
model.load_state_dict(
    pretrained_model.state_dict(),
    strict=False,
)

del pretrained_model

# Disable random patch augmentation on classification model
config = model.chronos_config

for attribute in [
    "masking_prob",
    "patch_masking_prob",
    "perturbation_prob",
    "patch_perturbation_prob",
]:
    if hasattr(config, attribute):
        setattr(config, attribute, 0.0)

# Fine-tune the entire model
for parameter in model.parameters():
    parameter.requires_grad = True

model.to(device)

Prepare the Dataset

ChronosBERT-Bolt processes numerical time series directly. The input consists of the time series context and a mask indicating which values are observed.

from torch.utils.data import Dataset, DataLoader

class ClassificationDataset(Dataset):

    def __init__(self, X, y, context_length):
        self.X = X
        self.y = np.asarray(y, dtype=np.int64)
        self.context_length = context_length

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

    def __getitem__(self, idx):

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

        mask = torch.isfinite(series)
        series[~mask] = float("nan")
        padding = self.context_length - len(series)

        if padding > 0:
            series = torch.cat([
                torch.full(
                    (padding,),
                    float("nan"),
                    dtype=series.dtype,
                ),
                series,
            ])

            mask = torch.cat([
                torch.zeros(padding, dtype=torch.bool),
                mask,
            ])

        return {
            "context": series,
            "mask": mask,
            "labels": torch.tensor(
                self.y[idx],
                dtype=torch.long,
            ),
        }


train_dataset = ClassificationDataset(
    X_train, y_train, context_length
)

test_dataset = ClassificationDataset(
    X_test, y_test, context_length
)

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

test_loader = DataLoader(
    test_dataset,
    batch_size=32,
)

4. Fine-tune the Model

The model is fine-tuned using cross-entropy loss and the AdamW optimizer. A linear learning rate scheduler with warmup is used during training.

import math
from torch.optim import AdamW
from transformers import get_linear_schedule_with_warmup

num_epochs = 10

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

for epoch in range(10):
    model.train()

    for batch in train_loader:
        context = batch["context"].to(device)
        mask = batch["mask"].to(device)
        labels = batch["labels"].to(device)

        outputs = model(
            context=context,
            mask=mask,
            labels=labels,
        )

        loss = outputs["loss"]
        loss.backward()

        optimizer.step()
        optimizer.zero_grad()

    print(f"Epoch {epoch + 1}, Loss: {loss.item():.4f}")

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}")