Simple Usage Overview

February 25, 2026 · View on GitHub

1. Train Prediction Heads

Train prediction heads on cached latent representations from a backbone model.

python src/train_heads.py \
    --model-name <MODEL> \
    --heads <HEAD_TYPE> \
    --batch-size 512 \
    --tot-iters 2097152 \
    --epochs 3 \
    --lr 3e-4 \
    --weight-decay 1e-2 \
    --grad-clip 1.0 \
    --use-amp \
    --save-dir models \
    --save-prefix <PREFIX> \
    --context-len 512 \
    --device cuda \
    --max-label 1000 \
    --min-label-default -1000
ArgumentDescription
--model-nameBackbone model: timesfm, chronos_bolt, moirai2, tirex, yinglong
--headsOne or more head types: mse, gaussian, poisson, neg_binom, studentst, quantiles, mixture
--mixture_componentsComponent distributions for mixture heads (one per head): moirai, gaussian, studentst
--mixture_KNumber of mixture components (one per head, integers)
--save-dirOutput directory for trained head checkpoints

2. Run Forecaster

Generate forecasts using a backbone model with trained heads.

Standard (non-AR) forecasting:

python src/run_forecaster.py \
    --dataset data/<DATASET>/y_<DATASET>.csv \
    --model-name <MODEL> \
    --save-dir results \
    --ckpt-dir models \
    --pred-length 64 \
    --context-len 512 \
    --forecast-date "2021-01-31 23:00:00" \
    --stride 1 \
    --batch-size 512 \
    --device cuda \
    --heads quantiles studentst \
    --mixture_K 3 3 \
    --mixture_components gaussian gaussian

Autoregressive forecasting:

python src/run_forecaster.py \
    --dataset data/<DATASET>/y_<DATASET>.csv \
    --model-name <MODEL> \
    --save-dir results \
    --ckpt-dir models \
    --pred-length 64 \
    --context-len 512 \
    --forecast-date "2021-01-31 23:00:00" \
    --stride 1 \
    --batch-size 512 \
    --device cuda \
    --heads \
    --ar \
    --ar-step-len 16 \
    --ar-samples 100
ArgumentDescription
--datasetPath to CSV with columns [unique_id, ds, y]
--headsTrained head names to evaluate; backbone is added automatically
--arEnable autoregressive forecasting
--ar-step-lenTokens generated per AR iteration (≤ model horizon)
--ar-samplesNumber of sampled trajectories for AR (omit for quantile-based AR)
--save-dirOutput directory for forecast CSVs
--ckpt-dirDirectory containing head checkpoints from step 1

3. Compute Metrics

Evaluate forecasts against ground-truth observations.

python src/metrics.py \
    --obs data/<DATASET>/y_<DATASET>.csv \
    --results_root results \
    --model_name <MODEL> \
    --head_name <HEAD> \
    --freq "H" \
    --q_low 0.1 \
    --q_high 0.9 \
    --confidence 0.8 \
    --wql_reduce mean \
    --wql_scale_mode sum_y \
    --out_per_h results/overall_per_h.csv \
    --out_per_h_by_series results/per_series_per_h.csv \
    --out_series_summary results/per_series_summary.csv \
    --out_summary results/overall_summary.json
ArgumentDescription
--obsPath to observations CSV with columns [unique_id, ds, y]
--results_rootRoot directory containing forecast outputs from step 2
--model_nameModel subdirectory under results_root
--head_nameHead subdirectory under model_name
--freqPandas frequency alias (D, H, 15min, W, etc.)
--out_per_hOutput path for per-horizon metrics CSV
--out_summaryOutput path for overall summary JSON
--out_per_h_by_seriesOutput path for per-series per-horizon metrics CSV
--out_series_summaryOutput path for per-series summary CSV

Metrics computed: MASE, SIW, PCE, TPCE, CCE, TCCE, WQL, MSIS — reported overall, per-horizon, and per-series.