| """ |
| Long Chronos-2 LoRA fine-tuning with large context window, tuned for local |
| CPU/memory. Single training run (Chronos does not support resuming LoRA fit). |
| |
| Usage (from project root): |
| PYTHONPATH=. python scripts/run_chronos_long_training.py [--device cpu] [--num-steps 4000] |
| |
| Uses: context_days=28, batch_size=16 by default. Set --num-steps for total steps. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| import sys |
| import threading |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parent.parent |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| |
| if "OMP_NUM_THREADS" not in os.environ: |
| try: |
| import multiprocessing |
| n = multiprocessing.cpu_count() |
| os.environ["OMP_NUM_THREADS"] = str(min(n, 10)) |
| except Exception: |
| pass |
|
|
| from config.settings import OUTPUTS_DIR |
| from src.chronos_forecaster import ( |
| ChronosForecaster, |
| STEPS_PER_DAY, |
| ) |
| from sklearn.metrics import mean_absolute_error |
| import numpy as np |
| import pandas as pd |
|
|
|
|
| def _quick_val_mae(forecaster: ChronosForecaster, df: pd.DataFrame, train_ratio: float, n_windows: int = 5) -> float: |
| """Compute MAE on first n_windows test windows (daytime-only) for convergence check.""" |
| sparse = forecaster.load_sparse_data() |
| daytime_ts = set(sparse["timestamp_utc"]) |
| split_idx = int(len(df) * train_ratio) |
| test_starts = list(range(split_idx, len(df) - STEPS_PER_DAY, STEPS_PER_DAY))[:n_windows] |
| actual_list, pred_list = [], [] |
| for start_idx in test_starts: |
| f = forecaster.forecast_day(df, start_idx, STEPS_PER_DAY, covariate_mode="all") |
| actual_slice = df.iloc[start_idx : start_idx + STEPS_PER_DAY] |
| daytime_mask = actual_slice["timestamp_utc"].isin(daytime_ts).values[:len(f)] |
| if daytime_mask.sum() < 5: |
| continue |
| actual_list.append(actual_slice["A"].values[:len(f)][daytime_mask]) |
| pred_list.append(np.clip(f["median"].values[daytime_mask], 0, None)) |
| if not actual_list: |
| return float("nan") |
| return float(mean_absolute_error(np.concatenate(actual_list), np.concatenate(pred_list))) |
|
|
|
|
| def main() -> None: |
| import argparse |
| parser = argparse.ArgumentParser(description="Chronos-2 long LoRA training (single run)") |
| parser.add_argument("--device", default="cpu", help="torch device (cpu or mps)") |
| parser.add_argument("--context-days", type=int, default=28, help="context window in days") |
| parser.add_argument("--batch-size", type=int, default=16, help="batch size (safe for 32GB RAM)") |
| parser.add_argument("--num-steps", type=int, default=4000, help="total training steps") |
| parser.add_argument("--learning-rate", type=float, default=1e-5, help="learning rate") |
| parser.add_argument("--progress-minutes", type=int, default=10, help="print timestamp and progress every N minutes") |
| parser.add_argument("--output-dir", type=str, default=None, help="output dir for checkpoints") |
| args = parser.parse_args() |
|
|
| output_dir = args.output_dir or str(OUTPUTS_DIR / "chronos_finetuned_long") |
| OUTPUTS_DIR.mkdir(parents=True, exist_ok=True) |
|
|
| print("Loading data...") |
| forecaster = ChronosForecaster(device=args.device, context_days=args.context_days) |
| df = forecaster.load_data() |
| print(f" Grid: {len(df)} rows, context={args.context_days}d ({forecaster.context_steps} steps)") |
| train_ratio = 0.75 |
| split_idx = int(len(df) * train_ratio) |
|
|
| print("\nBaseline (zero-shot) validation MAE (5 windows)...") |
| baseline_mae = _quick_val_mae(forecaster, df, train_ratio, n_windows=5) |
| print(f" {baseline_mae:.4f}") |
|
|
| print(f"\nLoRA fine-tuning: {args.num_steps} steps, batch_size={args.batch_size}, lr={args.learning_rate}...") |
| stop_event = threading.Event() |
| interval_sec = max(1, args.progress_minutes * 60) |
|
|
| def _progress_reporter(): |
| while True: |
| if stop_event.wait(interval_sec): |
| break |
| ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") |
| print(f"[{ts}] Chronos LoRA training still in progress ({args.num_steps} steps total)...", flush=True) |
|
|
| progress_thread = threading.Thread(target=_progress_reporter, daemon=True) |
| progress_thread.start() |
| try: |
| forecaster.finetune( |
| df, |
| train_ratio=train_ratio, |
| covariate_mode="all", |
| num_steps=args.num_steps, |
| learning_rate=args.learning_rate, |
| batch_size=args.batch_size, |
| output_dir=output_dir, |
| ) |
| finally: |
| stop_event.set() |
| progress_thread.join(timeout=interval_sec + 5) |
|
|
| print("\nValidation MAE after training (5 windows)...") |
| val_mae = _quick_val_mae(forecaster, df, train_ratio, n_windows=5) |
| print(f" {val_mae:.4f} (baseline {baseline_mae:.4f})") |
|
|
| |
| print("\nRunning full walk-forward benchmark (finetuned model, mode=all)...") |
| sparse = forecaster.load_sparse_data() |
| daytime_ts = set(sparse["timestamp_utc"]) |
| test_starts = list(range(split_idx, len(df) - STEPS_PER_DAY, STEPS_PER_DAY)) |
| all_actual, all_pred = [], [] |
| for start_idx in test_starts: |
| f = forecaster.forecast_day(df, start_idx, STEPS_PER_DAY, covariate_mode="all") |
| actual_slice = df.iloc[start_idx : start_idx + STEPS_PER_DAY] |
| daytime_mask = actual_slice["timestamp_utc"].isin(daytime_ts).values[:len(f)] |
| if daytime_mask.sum() < 5: |
| continue |
| all_actual.append(actual_slice["A"].values[:len(f)][daytime_mask]) |
| all_pred.append(np.clip(f["median"].values[daytime_mask], 0, None)) |
| lora_mae = None |
| if all_actual: |
| from sklearn.metrics import mean_squared_error, r2_score |
| a = np.concatenate(all_actual) |
| p = np.concatenate(all_pred) |
| lora_mae = float(mean_absolute_error(a, p)) |
| lora_rmse = float(np.sqrt(mean_squared_error(a, p))) |
| lora_r2 = float(r2_score(a, p)) |
| print(f" LoRA / all: MAE={lora_mae:.4f} RMSE={lora_rmse:.4f} R²={lora_r2:.4f} ({len(all_actual)} windows, {len(a)} steps)") |
|
|
| |
| bench_path = OUTPUTS_DIR / "chronos_benchmark.csv" |
| if bench_path.exists(): |
| existing = pd.read_csv(bench_path) |
| lora_row = pd.DataFrame([{ |
| "mode": "lora / all", |
| "MAE": lora_mae, |
| "RMSE": lora_rmse, |
| "R2": lora_r2, |
| "n_windows": len(all_actual), |
| "n_steps": len(a), |
| }]) |
| combined = pd.concat([existing, lora_row], ignore_index=True) |
| combined.to_csv(bench_path, index=False) |
| print(f" Appended lora row → {bench_path}") |
| else: |
| pd.DataFrame([{ |
| "mode": "lora / all", "MAE": lora_mae, "RMSE": lora_rmse, "R2": lora_r2, |
| "n_windows": len(all_actual), "n_steps": len(a), |
| }]).to_csv(bench_path, index=False) |
|
|
| |
| print("\nGenerating sample forecast plot...") |
| forecaster.plot_sample_forecast(df) |
|
|
| |
| print("\n" + "=" * 60) |
| print("TRAINING COMPLETE — Next steps") |
| print("=" * 60) |
| print(f" Checkpoints: {output_dir}") |
| print(f" Benchmark: {OUTPUTS_DIR / 'chronos_benchmark.csv'} (lora / all row appended)") |
| print(f" Plot: {OUTPUTS_DIR / 'chronos_forecast_sample.png'} (from finetuned model)") |
| print(" • Refresh the Streamlit app: Prediction Accuracy tab will show LoRA / all in the table.") |
| print(" • Sample forecast image is from the finetuned model.") |
| if lora_mae is not None: |
| zs_mae = 3.91 |
| delta = (zs_mae - lora_mae) / zs_mae * 100 |
| print(f" • LoRA MAE {lora_mae:.2f} vs zero-shot ~{zs_mae:.2f} ({delta:+.0f}% change).") |
| print("=" * 60) |
| print("Done.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
