Refactor into 2 examples
This commit is contained in:
misha-chertushkin
+153
-228
@@ -1,274 +1,199 @@
|
||||
# Filename: tutorial_timesfm.py
|
||||
"""
|
||||
TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
|
||||
|
||||
Example usage:
|
||||
```python
|
||||
# Prepare datasets
|
||||
train_dataset = TimeSeriesDataset(train_data, context_length=128, horizon_length=32)
|
||||
val_dataset = TimeSeriesDataset(val_data, context_length=128, horizon_length=32)
|
||||
|
||||
# Initialize model and configuration
|
||||
model = TimesFm(...)
|
||||
config = FinetuningConfig(
|
||||
batch_size=64,
|
||||
num_epochs=50,
|
||||
learning_rate=1e-4,
|
||||
use_wandb=True
|
||||
)
|
||||
|
||||
# Create finetuner
|
||||
finetuner = TimesFMFinetuner(model, config)
|
||||
|
||||
# Finetune model
|
||||
results = finetuner.finetune(train_dataset, val_dataset)
|
||||
```
|
||||
"""
|
||||
|
||||
import abc
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
import yfinance as yf
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
import torch.optim as optim
|
||||
import timesfm
|
||||
from os import path
|
||||
from typing import Any, Sequence
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
import wandb
|
||||
import multiprocessing as mp
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from timesfm import TimesFm
|
||||
|
||||
|
||||
from timesfm.pytorch_patched_decoder import TimesFMConfig, PatchedTimeSeriesDecoder
|
||||
@dataclass
|
||||
class FinetuningConfig:
|
||||
"""Configuration for TimesFM finetuning process."""
|
||||
|
||||
import torch
|
||||
import matplotlib.pyplot as plt
|
||||
# Training parameters
|
||||
batch_size: int = 32
|
||||
num_epochs: int = 20
|
||||
learning_rate: float = 1e-4
|
||||
weight_decay: float = 0.01
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
# Hardware parameters
|
||||
device: str = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
distributed: bool = False
|
||||
world_size: int = 1
|
||||
|
||||
# Logging parameters
|
||||
use_wandb: bool = False
|
||||
wandb_project: str = "timesfm-finetuning"
|
||||
|
||||
|
||||
# --------------------------------------------------
|
||||
# 1. Download stock data via yfinance
|
||||
# --------------------------------------------------
|
||||
def download_yfinance_data(ticker="AAPL", start="2020-01-01", end="2022-01-01"):
|
||||
"""
|
||||
Download daily stock data for a given ticker from Yahoo Finance.
|
||||
Returns a pandas DataFrame with columns like 'Open', 'High', 'Low', 'Close', 'Volume'.
|
||||
"""
|
||||
df = yf.download(ticker, start=start, end=end)
|
||||
df = df.dropna()
|
||||
return df["Close"].reset_index(drop=True)
|
||||
class TimesFMFinetuner:
|
||||
"""Main class for finetuning TimesFM models."""
|
||||
|
||||
|
||||
# --------------------------------------------------
|
||||
# 2. Create a dataset class for TimesFM
|
||||
# --------------------------------------------------
|
||||
class FinancialDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
series: pd.Series,
|
||||
config: TimesFMConfig,
|
||||
context_length=128, # how many past timesteps as input
|
||||
horizon_length=32, # how many future steps to predict
|
||||
model: TimesFm,
|
||||
config: FinetuningConfig,
|
||||
loss_fn: Optional[callable] = None,
|
||||
logger: Optional[logging.Logger] = None,
|
||||
):
|
||||
super().__init__()
|
||||
"""
|
||||
Initialize TimesFM finetuner.
|
||||
|
||||
self.series = series.values.astype(np.float32)
|
||||
self.context_length = context_length
|
||||
self.horizon_length = horizon_length
|
||||
Args:
|
||||
model: TimesFM model to finetune
|
||||
config: Finetuning configuration
|
||||
logger: Optional logger instance
|
||||
"""
|
||||
self.model = model
|
||||
self.config = config
|
||||
self.logger = logger or logging.getLogger(__name__)
|
||||
|
||||
self.samples = []
|
||||
# We want to ensure we have at least context_length + horizon_length points.
|
||||
for start_idx in range(0, len(self.series) - (context_length + horizon_length)):
|
||||
end_idx = start_idx + context_length
|
||||
# context slice
|
||||
x_context = self.series[start_idx:end_idx]
|
||||
# future/horizon slice
|
||||
x_future = self.series[end_idx : end_idx + horizon_length]
|
||||
self.samples.append((x_context, x_future))
|
||||
self.device = torch.device(config.device)
|
||||
self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2)) # MSELoss()
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
if config.use_wandb:
|
||||
self._setup_wandb()
|
||||
|
||||
def __getitem__(self, index):
|
||||
x_context, x_future = self.samples[index]
|
||||
# Convert to torch
|
||||
x_context = torch.tensor(x_context, dtype=torch.float32)
|
||||
x_future = torch.tensor(x_future, dtype=torch.float32)
|
||||
def _setup_wandb(self) -> None:
|
||||
"""Initialize Weights & Biases logging."""
|
||||
wandb.init(project=self.config.wandb_project, entity=self.config.wandb_entity, config=self.config.__dict__)
|
||||
|
||||
input_padding = torch.zeros_like(x_context)
|
||||
def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader:
|
||||
"""Create a dataloader from a dataset."""
|
||||
return DataLoader(
|
||||
dataset,
|
||||
batch_size=self.config.batch_size,
|
||||
shuffle=name == "train",
|
||||
num_workers=mp.cpu_count(),
|
||||
pin_memory=self.device.type == "cuda",
|
||||
persistent_workers=True,
|
||||
prefetch_factor=2,
|
||||
)
|
||||
|
||||
freq = torch.zeros(1, dtype=torch.long)
|
||||
def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
|
||||
"""Train for one epoch."""
|
||||
self.model.train()
|
||||
total_loss = 0.0
|
||||
n_batches = len(train_loader)
|
||||
|
||||
return x_context, input_padding, freq, x_future
|
||||
for batch in train_loader:
|
||||
x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
|
||||
|
||||
|
||||
def collate_fn(batch):
|
||||
xs_context = [item[0] for item in batch]
|
||||
xs_padding = [item[1] for item in batch]
|
||||
freqs = [item[2] for item in batch]
|
||||
xs_future = [item[3] for item in batch]
|
||||
|
||||
x_context = torch.stack(xs_context, dim=0)
|
||||
input_pad = torch.stack(xs_padding, dim=0)
|
||||
freq = torch.stack(freqs, dim=0) # shape [B, 1]
|
||||
x_future = torch.stack(xs_future, dim=0)
|
||||
|
||||
return x_context, input_pad, freq, x_future
|
||||
|
||||
|
||||
def get_model(*, load_weights: bool = False):
|
||||
# standard model hack
|
||||
repo_id = "google/timesfm-2.0-500m-pytorch"
|
||||
tfm = timesfm.TimesFm(
|
||||
hparams=timesfm.TimesFmHparams(
|
||||
backend="cuda",
|
||||
per_core_batch_size=32,
|
||||
horizon_len=128,
|
||||
num_layers=50,
|
||||
use_positional_embedding=False,
|
||||
context_len=192,
|
||||
),
|
||||
checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=repo_id),
|
||||
)
|
||||
|
||||
model = PatchedTimeSeriesDecoder(tfm._model_config)
|
||||
|
||||
if load_weights:
|
||||
checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
|
||||
print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
|
||||
loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
|
||||
model.load_state_dict(loaded_checkpoint)
|
||||
print("After loading:")
|
||||
print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
|
||||
model = model.to(device)
|
||||
|
||||
# import sys
|
||||
# sys.exit(-1)
|
||||
# repo_id = "google/timesfm-1.0-200m"
|
||||
return model, tfm._model_config
|
||||
|
||||
|
||||
def train_model(
|
||||
ticker="AAPL", start="2015-01-01", end="2022-01-01", train_split=0.8, batch_size=8, num_epochs=20, pretrained=False
|
||||
):
|
||||
df_close = download_yfinance_data(ticker, start=start, end=end)
|
||||
model, config = get_model(load_weights=pretrained)
|
||||
|
||||
total_len = len(df_close)
|
||||
train_size = int(total_len * train_split)
|
||||
val_size = total_len - train_size
|
||||
|
||||
train_series = df_close.iloc[:train_size].reset_index(drop=True)
|
||||
val_series = df_close.iloc[train_size:].reset_index(drop=True)
|
||||
|
||||
train_dataset = FinancialDataset(
|
||||
series=train_series, config=config, context_length=128, horizon_length=config.horizon_len
|
||||
)
|
||||
val_dataset = FinancialDataset(
|
||||
series=val_series, config=config, context_length=128, horizon_length=config.horizon_len
|
||||
)
|
||||
print("Train samples:", len(train_dataset))
|
||||
print("Val samples:", len(val_dataset))
|
||||
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
|
||||
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
|
||||
|
||||
optimizer = optim.Adam(model.parameters(), lr=1e-4)
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
model.train()
|
||||
total_train_loss = 0.0
|
||||
|
||||
for x_context, x_padding, freq, x_future in train_dataloader:
|
||||
x_context, x_padding, freq, x_future = (
|
||||
x_context.to(device),
|
||||
x_padding.to(device),
|
||||
freq.to(device),
|
||||
x_future.to(device),
|
||||
)
|
||||
predictions = model(x_context, x_padding.float(), freq)
|
||||
# predictions shape => [B, N, horizon_len, (1 + #quantiles)]
|
||||
predictions_mean = predictions[..., 0] # => [B, N, horizon_len]
|
||||
last_patch_pred = predictions_mean[:, -1, :] # => [B, horizon_len]
|
||||
|
||||
# x_future => [B, horizon_len]
|
||||
loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
|
||||
predictions = self.model(x_context, x_padding.float(), freq)
|
||||
predictions_mean = predictions[..., 0]
|
||||
last_patch_pred = predictions_mean[:, -1, :]
|
||||
loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
total_train_loss += loss.item()
|
||||
total_loss += loss.item()
|
||||
|
||||
avg_train_loss = total_train_loss / len(train_dataloader)
|
||||
return total_loss / n_batches
|
||||
|
||||
# -------- Compute validation loss --------
|
||||
model.eval()
|
||||
total_val_loss = 0.0
|
||||
with torch.no_grad():
|
||||
for x_context, x_padding, freq, x_future in val_dataloader:
|
||||
x_context, x_padding, freq, x_future = (
|
||||
x_context.to(device),
|
||||
x_padding.to(device),
|
||||
freq.to(device),
|
||||
x_future.to(device),
|
||||
)
|
||||
predictions = model(x_context, x_padding.float(), freq)
|
||||
predictions_mean = predictions[..., 0]
|
||||
last_patch_pred = predictions_mean[:, -1, :]
|
||||
val_loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
|
||||
total_val_loss += val_loss.item()
|
||||
@torch.no_grad()
|
||||
def _validate(self, val_loader: DataLoader) -> float:
|
||||
"""Perform validation."""
|
||||
self.model.eval()
|
||||
total_loss = 0.0
|
||||
|
||||
avg_val_loss = total_val_loss / max(len(val_dataloader), 1)
|
||||
for batch in val_loader:
|
||||
x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch]
|
||||
|
||||
print(f"[Epoch {epoch+1}] Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
|
||||
predictions = self.model(x_context, x_padding.float(), freq)
|
||||
predictions_mean = predictions[..., 0]
|
||||
last_patch_pred = predictions_mean[:, -1, :]
|
||||
|
||||
torch.save(model.state_dict(), "timesfm_finetuned.pth")
|
||||
return model, train_dataloader, val_dataloader
|
||||
loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
|
||||
total_loss += loss.item()
|
||||
|
||||
return total_loss / len(val_loader)
|
||||
|
||||
def plot_predictions(model, dataloader):
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
x_context, x_padding, freq, x_future = next(iter(dataloader))
|
||||
x_context, x_padding, freq, x_future = (
|
||||
x_context.to(device),
|
||||
x_padding.to(device),
|
||||
freq.to(device),
|
||||
x_future.to(device),
|
||||
)
|
||||
# Forward pass
|
||||
predictions = model(x_context, x_padding.float(), freq)
|
||||
# => [B, N, horizon_len, (1 + #quantiles)]
|
||||
predictions_mean = predictions[..., 0] # => [B, N, horizon_len]
|
||||
last_patch_prediction = predictions_mean[:, -1, :] # => [B, horizon_len]
|
||||
def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
|
||||
"""
|
||||
Finetune the TimesFM model on the provided datasets.
|
||||
|
||||
# We'll plot only the first sample in the batch
|
||||
i = 0
|
||||
pred_vals = last_patch_prediction[i].cpu().numpy() # [horizon_len]
|
||||
context_vals = x_context[i].cpu().numpy() # [context_len]
|
||||
future_vals = x_future[i].cpu().numpy() # [horizon_len]
|
||||
Args:
|
||||
train_dataset: Training dataset
|
||||
val_dataset: Validation dataset
|
||||
|
||||
horizon_len = future_vals.shape[0]
|
||||
context_len = context_vals.shape[0]
|
||||
Returns:
|
||||
Dict containing training history and best model path
|
||||
"""
|
||||
self.model = self.model.to(self.device)
|
||||
|
||||
plt.figure(figsize=(10, 5))
|
||||
train_loader = self._create_dataloader(train_dataset, "train")
|
||||
val_loader = self._create_dataloader(val_dataset, "val")
|
||||
|
||||
# Plot context
|
||||
plt.plot(range(context_len), context_vals, label="Context (History)", color="blue")
|
||||
|
||||
# Plot predicted future
|
||||
plt.plot(
|
||||
range(context_len, context_len + horizon_len),
|
||||
pred_vals,
|
||||
label="Predicted Future",
|
||||
color="orange",
|
||||
optimizer = optim.Adam(
|
||||
self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
|
||||
)
|
||||
|
||||
# Plot ground truth future
|
||||
plt.plot(
|
||||
range(context_len, context_len + horizon_len),
|
||||
future_vals,
|
||||
label="Ground Truth Future",
|
||||
color="green",
|
||||
linestyle="--",
|
||||
)
|
||||
history = {"train_loss": [], "val_loss": [], "learning_rate": []}
|
||||
|
||||
plt.xlabel("Time")
|
||||
plt.ylabel("Value")
|
||||
plt.title("Model Forecast vs. Ground Truth")
|
||||
plt.legend()
|
||||
plt.show()
|
||||
plt.savefig("pic_predictions.png")
|
||||
self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
|
||||
self.logger.info(f"Training samples: {len(train_dataset)}")
|
||||
self.logger.info(f"Validation samples: {len(val_dataset)}")
|
||||
|
||||
try:
|
||||
for epoch in range(self.config.num_epochs):
|
||||
train_loss = self._train_epoch(train_loader, optimizer)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage
|
||||
model, train_dl, val_dl = train_model(
|
||||
ticker="AAPL",
|
||||
start="2012-01-01",
|
||||
end="2019-01-01",
|
||||
train_split=0.8,
|
||||
batch_size=256,
|
||||
num_epochs=50,
|
||||
pretrained=True,
|
||||
)
|
||||
val_loss = self._validate(val_loader)
|
||||
|
||||
plot_predictions(model, val_dl)
|
||||
current_lr = optimizer.param_groups[0]["lr"]
|
||||
|
||||
history["train_loss"].append(train_loss)
|
||||
history["val_loss"].append(val_loss)
|
||||
history["learning_rate"].append(current_lr)
|
||||
|
||||
metrics = {
|
||||
"train_loss": train_loss,
|
||||
"val_loss": val_loss,
|
||||
"learning_rate": current_lr,
|
||||
"epoch": epoch + 1,
|
||||
}
|
||||
|
||||
if self.config.use_wandb:
|
||||
wandb.log(metrics)
|
||||
|
||||
print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
|
||||
except KeyboardInterrupt:
|
||||
self.logger.info("Training interrupted by user")
|
||||
|
||||
return {"history": history}
|
||||
|
||||
Reference in New Issue
Block a user