misha-chertushkin
341 lines
11 KiB
Python
341 lines
11 KiB
Python
"""
|
|
TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
import wandb
|
|
|
|
|
|
class MetricsLogger(ABC):
|
|
"""Abstract base class for logging metrics during training.
|
|
|
|
This class defines the interface for logging metrics during model training.
|
|
Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
|
|
"""
|
|
|
|
@abstractmethod
|
|
def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
|
|
"""Log metrics to the specified backend.
|
|
|
|
Args:
|
|
metrics: Dictionary containing metric names and values.
|
|
step: Optional step number or epoch for the metrics.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def close(self) -> None:
|
|
"""Clean up any resources used by the logger."""
|
|
pass
|
|
|
|
|
|
class WandBLogger(MetricsLogger):
|
|
"""Weights & Biases implementation of metrics logging.
|
|
|
|
Args:
|
|
project: Name of the W&B project.
|
|
config: Configuration dictionary to log.
|
|
rank: Process rank in distributed training.
|
|
"""
|
|
|
|
def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
|
|
self.rank = rank
|
|
if rank == 0:
|
|
wandb.init(project=project, config=config)
|
|
|
|
def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
|
|
"""Log metrics to W&B if on the main process.
|
|
|
|
Args:
|
|
metrics: Dictionary of metrics to log.
|
|
step: Current training step or epoch.
|
|
"""
|
|
if self.rank == 0:
|
|
wandb.log(metrics, step=step)
|
|
|
|
def close(self) -> None:
|
|
"""Finish the W&B run if on the main process."""
|
|
if self.rank == 0:
|
|
wandb.finish()
|
|
|
|
|
|
class DistributedManager:
|
|
"""Manages distributed training setup and cleanup.
|
|
|
|
Args:
|
|
world_size: Total number of processes.
|
|
rank: Process rank.
|
|
master_addr: Address of the master process.
|
|
master_port: Port for distributed communication.
|
|
backend: PyTorch distributed backend to use.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
world_size: int,
|
|
rank: int,
|
|
master_addr: str = "localhost",
|
|
master_port: str = "12358",
|
|
backend: str = "nccl",
|
|
):
|
|
self.world_size = world_size
|
|
self.rank = rank
|
|
self.master_addr = master_addr
|
|
self.master_port = master_port
|
|
self.backend = backend
|
|
|
|
def setup(self) -> None:
|
|
"""Initialize the distributed environment."""
|
|
os.environ["MASTER_ADDR"] = self.master_addr
|
|
os.environ["MASTER_PORT"] = self.master_port
|
|
|
|
if not dist.is_initialized():
|
|
dist.init_process_group(backend=self.backend, world_size=self.world_size, rank=self.rank)
|
|
|
|
def cleanup(self) -> None:
|
|
"""Clean up the distributed environment."""
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@dataclass
|
|
class FinetuningConfig:
|
|
"""Configuration for model training.
|
|
|
|
Args:
|
|
batch_size: Number of samples per batch.
|
|
num_epochs: Number of training epochs.
|
|
learning_rate: Initial learning rate.
|
|
weight_decay: L2 regularization factor.
|
|
device: Device to train on ('cuda' or 'cpu').
|
|
distributed: Whether to use distributed training.
|
|
gpu_ids: List of GPU IDs to use.
|
|
master_port: Port for distributed training.
|
|
master_addr: Address for distributed training.
|
|
use_wandb: Whether to use Weights & Biases logging.
|
|
wandb_project: W&B project name.
|
|
"""
|
|
|
|
batch_size: int = 32
|
|
num_epochs: int = 20
|
|
learning_rate: float = 1e-4
|
|
weight_decay: float = 0.01
|
|
device: str = "cuda" if torch.cuda.is_available() else "cpu"
|
|
distributed: bool = False
|
|
gpu_ids: List[int] = field(default_factory=lambda: [0])
|
|
master_port: str = "12358"
|
|
master_addr: str = "localhost"
|
|
use_wandb: bool = False
|
|
wandb_project: str = "timesfm-finetuning"
|
|
|
|
|
|
class TimesFMFinetuner:
|
|
"""Handles model training and validation.
|
|
|
|
Args:
|
|
model: PyTorch model to train.
|
|
config: Training configuration.
|
|
rank: Process rank for distributed training.
|
|
loss_fn: Loss function (defaults to MSE).
|
|
logger: Optional logging.Logger instance.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: nn.Module,
|
|
config: FinetuningConfig,
|
|
rank: int = 0,
|
|
loss_fn: Optional[Callable] = None,
|
|
logger: Optional[logging.Logger] = None,
|
|
):
|
|
self.model = model
|
|
self.config = config
|
|
self.rank = rank
|
|
self.logger = logger or logging.getLogger(__name__)
|
|
self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
|
|
self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))
|
|
|
|
if config.use_wandb:
|
|
self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__, rank)
|
|
|
|
if config.distributed:
|
|
self.dist_manager = DistributedManager(
|
|
world_size=len(config.gpu_ids),
|
|
rank=rank,
|
|
master_addr=config.master_addr,
|
|
master_port=config.master_port,
|
|
)
|
|
self.dist_manager.setup()
|
|
self.model = self._setup_distributed_model()
|
|
|
|
def _setup_distributed_model(self) -> nn.Module:
|
|
"""Configure model for distributed training."""
|
|
self.model = self.model.to(self.device)
|
|
return DDP(
|
|
self.model, device_ids=[self.config.gpu_ids[self.rank]], output_device=self.config.gpu_ids[self.rank]
|
|
)
|
|
|
|
def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
|
|
"""Create appropriate DataLoader based on training configuration.
|
|
|
|
Args:
|
|
dataset: Dataset to create loader for.
|
|
is_train: Whether this is for training (affects shuffling).
|
|
|
|
Returns:
|
|
DataLoader instance.
|
|
"""
|
|
if self.config.distributed:
|
|
sampler = torch.utils.data.distributed.DistributedSampler(
|
|
dataset, num_replicas=len(self.config.gpu_ids), rank=dist.get_rank(), shuffle=is_train
|
|
)
|
|
else:
|
|
sampler = None
|
|
|
|
return DataLoader(
|
|
dataset,
|
|
batch_size=self.config.batch_size,
|
|
shuffle=(is_train and not self.config.distributed),
|
|
sampler=sampler,
|
|
)
|
|
|
|
def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
|
|
"""Process a single batch of data.
|
|
|
|
Args:
|
|
batch: List of input tensors.
|
|
|
|
Returns:
|
|
Tuple of (loss, predictions).
|
|
"""
|
|
x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
|
|
|
|
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))
|
|
|
|
return loss, predictions
|
|
|
|
def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
|
|
"""Train for one epoch.
|
|
|
|
Args:
|
|
train_loader: DataLoader for training data.
|
|
optimizer: Optimizer instance.
|
|
|
|
Returns:
|
|
Average training loss for the epoch.
|
|
"""
|
|
self.model.train()
|
|
total_loss = 0.0
|
|
|
|
for batch in train_loader:
|
|
loss, _ = self._process_batch(batch)
|
|
|
|
if self.config.distributed:
|
|
losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
|
|
dist.all_gather(losses, loss)
|
|
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
total_loss += loss.item()
|
|
|
|
return total_loss / len(train_loader)
|
|
|
|
def _validate(self, val_loader: DataLoader) -> float:
|
|
"""Perform validation.
|
|
|
|
Args:
|
|
val_loader: DataLoader for validation data.
|
|
|
|
Returns:
|
|
Average validation loss.
|
|
"""
|
|
self.model.eval()
|
|
total_loss = 0.0
|
|
|
|
with torch.no_grad():
|
|
for batch in val_loader:
|
|
loss, _ = self._process_batch(batch)
|
|
|
|
if self.config.distributed:
|
|
losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
|
|
dist.all_gather(losses, loss)
|
|
|
|
total_loss += loss.item()
|
|
|
|
return total_loss / len(val_loader)
|
|
|
|
def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
|
|
"""Train the model.
|
|
|
|
Args:
|
|
train_dataset: Training dataset.
|
|
val_dataset: Validation dataset.
|
|
|
|
Returns:
|
|
Dictionary containing training history.
|
|
"""
|
|
self.model = self.model.to(self.device)
|
|
train_loader = self._create_dataloader(train_dataset, is_train=True)
|
|
val_loader = self._create_dataloader(val_dataset, is_train=False)
|
|
|
|
optimizer = torch.optim.Adam(
|
|
self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
|
|
)
|
|
|
|
history = {"train_loss": [], "val_loss": [], "learning_rate": []}
|
|
|
|
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)
|
|
val_loss = self._validate(val_loader)
|
|
current_lr = optimizer.param_groups[0]["lr"]
|
|
|
|
metrics = {
|
|
"train_loss": train_loss,
|
|
"val_loss": val_loss,
|
|
"learning_rate": current_lr,
|
|
"epoch": epoch + 1,
|
|
}
|
|
|
|
if self.config.use_wandb:
|
|
self.metrics_logger.log_metrics(metrics)
|
|
|
|
history["train_loss"].append(train_loss)
|
|
history["val_loss"].append(val_loss)
|
|
history["learning_rate"].append(current_lr)
|
|
|
|
if self.rank == 0:
|
|
self.logger.info(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
|
|
|
|
except KeyboardInterrupt:
|
|
self.logger.info("Training interrupted by user")
|
|
|
|
if self.config.distributed:
|
|
self.dist_manager.cleanup()
|
|
|
|
if self.config.use_wandb:
|
|
self.metrics_logger.close()
|
|
|
|
return {"history": history}
|