training/fine-tuning clarification

[Mhdaw]

Hello, I want to train/fine-tune the chronos models, Is this a good approach?

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForSeq2SeqLM, AdamW
import numpy as np
from tqdm import tqdm
from chronos import ChronosConfig, MeanScaleUniformBins, ChronosModel

config = ChronosConfig(
    tokenizer_class="MeanScaleUniformBins",
    tokenizer_kwargs={"low_limit": -5.0, "high_limit": 10.0},  # Not sure about how these values should be!
    context_length=20,  
    prediction_length=10, 
    n_tokens=100,       # 4096
    n_special_tokens=2, # padding and eos?
    pad_token_id=0,
    eos_token_id=1,
    use_eos_token=False,
    model_type="seq2seq", 
    num_samples=10,
    temperature=1.0,
    top_k=50,
    top_p=0.95,
)
tokenizer = MeanScaleUniformBins(
    low_limit=config.tokenizer_kwargs["low_limit"],
    high_limit=config.tokenizer_kwargs["high_limit"],
    config=config
)
base_model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-efficient-tiny", device_map="auto")  # Or another T5 variant like amazon/chronos-t5-tiny

model = ChronosModel(config=config, model=base_model)

class TimeSeriesDataset(Dataset):
    def __init__(self, data, context_length, prediction_length):
        self.data = data
        self.context_length = context_length
        self.prediction_length = prediction_length

    def __len__(self):
        return len(self.data) - self.context_length - self.prediction_length + 1

    def __getitem__(self, idx):
        context = self.data[idx : idx + self.context_length]
        target = self.data[idx + self.context_length : idx + self.context_length + self.prediction_length]
        return torch.tensor(context), torch.tensor(target)

raw_data = df["target"].values
train_size = int(0.8 * len(raw_data))
train_data = raw_data[:train_size]
val_data = raw_data[train_size:]

train_dataset = TimeSeriesDataset(train_data, config.context_length, config.prediction_length)
val_dataset = TimeSeriesDataset(val_data, config.context_length, config.prediction_length)
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32)

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

num_epochs = 5
for epoch in tqdm(range(num_epochs), desc="Epochs"):  
    model.train()
    total_loss = 0
    for batch in train_dataloader:
        context, target = batch
        context = context.to(model.device)
        target = target.to(model.device)

        tokenizer.boundaries = tokenizer.boundaries.to(model.device)  
        context_tokens, context_attention_mask, scale = tokenizer.context_input_transform(context)

        target_tokens, target_attention_mask = tokenizer.label_input_transform(target, scale)

        outputs = model.model(
            input_ids=context_tokens,
            attention_mask=context_attention_mask,
            labels=target_tokens,  
            decoder_attention_mask=target_attention_mask, 
        )

        loss = outputs.loss
        total_loss += loss.item()

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

    avg_loss = total_loss / len(train_dataloader)
    tqdm.write(f"Epoch {epoch+1}/{num_epochs}, Training Loss: {avg_loss:.4f}") 

    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch in val_dataloader:
            context, target = batch
            context = context.to(model.device)
            target = target.to(model.device)
            tokenizer.boundaries = tokenizer.boundaries.to(model.device)  
            context_tokens, context_attention_mask, scale = tokenizer.context_input_transform(context)
            target_tokens, target_attention_mask = tokenizer.label_input_transform(target, scale)
            outputs = model.model(
                input_ids=context_tokens,
                attention_mask=context_attention_mask,
                labels=target_tokens,
                decoder_attention_mask=target_attention_mask,
            )
            val_loss += outputs.loss.item()
    avg_val_loss = val_loss / len(val_dataloader)
    tqdm.write(f"Epoch {epoch+1}/{num_epochs}, Validation Loss: {avg_val_loss:.4f}")

model.model.save_pretrained("chronos-t5-timeseries")

is there any thing that I've missed? And any advice on the tokenizer_kwargs and n_special_tokens

[Killer3048]

@Mhdaw I think your approach is on the right track! Here are a few suggestions to enhance your training:

Training Tips

• Set use_eos_token=True for seq2seq models to signify the end of sequences.
• Optimize the learning rate (e.g., try 3e-5) and consider increasing epochs for better convergence.
• Incorporate additional validation metrics like MAE or RMSE to monitor performance.

[Mhdaw]

Thanks, I think the use_eos_token=True whould help with models performance.
adding metrics like MAE or RMSE requiers few steps, one that this model outputs are classes(in train phase) After that it outputs a probability distribution, maybe using the median and target for the metric calculation could be the answer.

[Mhdaw]

Also, If I want to use the Chronos Bolt models, What changes I need to make?

[Killer3048]

You can check how they implemented it themselves chronos-bolt.py