Sync from GitHub (tests passed)

deep_learning/training/hyperopt.py

@@ -43,26 +43,31 @@ def create_trial_config(trial, base_cfg: TFTASROConfig) -> TFTASROConfig:
     model_cfg = TFTModelConfig(
         max_encoder_length=trial.suggest_int("max_encoder_length", 30, 90, step=10),
         max_prediction_length=base_cfg.model.max_prediction_length,
-        hidden_size=trial.suggest_int("hidden_size", 32, 128, step=16),
-        attention_head_size=trial.suggest_int("attention_head_size", 1, 8),
-        dropout=trial.suggest_float("dropout", 0.05, 0.3, step=0.05),
-        hidden_continuous_size=trial.suggest_int("hidden_continuous_size", 16, 64, step=8),
+        # Cap at 64: beyond that the VSN encoder explodes to 3M+ params for our
+        # 313-sample dataset, causing the same overfitting we already saw at 64.
+        hidden_size=trial.suggest_int("hidden_size", 16, 64, step=16),
+        attention_head_size=trial.suggest_int("attention_head_size", 1, 4),
+        dropout=trial.suggest_float("dropout", 0.1, 0.5, step=0.05),
+        hidden_continuous_size=trial.suggest_int("hidden_continuous_size", 8, 32, step=8),
         quantiles=base_cfg.model.quantiles,
-        learning_rate=trial.suggest_float("learning_rate", 1e-4, 1e-2, log=True),
+        learning_rate=trial.suggest_float("learning_rate", 5e-5, 5e-3, log=True),
         reduce_on_plateau_patience=4,
-        gradient_clip_val=trial.suggest_float("gradient_clip_val", 0.1, 1.0, step=0.1),
+        gradient_clip_val=trial.suggest_float("gradient_clip_val", 0.5, 2.0, step=0.5),
     )
 
     asro_cfg = ASROConfig(
-        lambda_vol=trial.suggest_float("lambda_vol", 0.1, 0.5, step=0.05),
-        lambda_quantile=trial.suggest_float("lambda_quantile", 0.1, 0.5, step=0.05),
+        lambda_vol=trial.suggest_float("lambda_vol", 0.1, 0.4, step=0.05),
+        # lambda_quantile is the explicit w_quantile weight (w_sharpe = 1 - w_q)
+        lambda_quantile=trial.suggest_float("lambda_quantile", 0.2, 0.6, step=0.05),
         risk_free_rate=0.0,
     )
 
     training_cfg = TrainingConfig(
         max_epochs=50,
         early_stopping_patience=8,
-        batch_size=trial.suggest_categorical("batch_size", [32, 64, 128]),
+        # Include 16 which gives 19 batches/epoch (vs 4 at batch_size=64)
+        # — more gradient steps per epoch → more stable convergence.
+        batch_size=trial.suggest_categorical("batch_size", [16, 32, 64]),
         val_ratio=base_cfg.training.val_ratio,
         test_ratio=base_cfg.training.test_ratio,
         lookback_days=base_cfg.training.lookback_days,

deep_learning/training/trainer.py

@@ -68,7 +68,13 @@ def train_tft_model(
     if cfg is None:
         cfg = get_tft_config()
 
-    # ---- 0. ASRO loss sanity check (runs before any training) ----
+    # ---- 0a. Load Optuna best params if available ----
+    # When the hyperopt step ran before this trainer, it writes best params to
+    # optuna_results.json. We apply those params over the default config so that
+    # the final training run actually benefits from the search.
+    cfg = _apply_optuna_results(cfg)
+
+    # ---- 0b. ASRO loss sanity check (runs before any training) ----
     try:
         from deep_learning.models.losses import debug_asro_loss_direction
         debug = debug_asro_loss_direction()
@@ -270,6 +276,58 @@ def train_tft_model(
     return result
 
 
+def _apply_optuna_results(cfg: TFTASROConfig) -> TFTASROConfig:
+    """
+    If an optuna_results.json exists in the checkpoint directory, overlay the
+    best hyperparameters onto cfg and return the updated config.  This connects
+    the hyperopt step to the final training run so search results are not wasted.
+    """
+    import json
+    from dataclasses import replace
+    from deep_learning.config import ASROConfig, TFTModelConfig, TrainingConfig
+
+    results_path = Path(cfg.training.checkpoint_dir) / "optuna_results.json"
+    if not results_path.exists():
+        return cfg
+
+    try:
+        data = json.loads(results_path.read_text())
+        params = data.get("best_params", {})
+        if not params:
+            return cfg
+
+        model_overrides = {
+            k: params[k] for k in (
+                "hidden_size", "attention_head_size", "dropout",
+                "hidden_continuous_size", "learning_rate",
+                "gradient_clip_val", "max_encoder_length",
+            ) if k in params
+        }
+        asro_overrides = {
+            k: params[k] for k in ("lambda_vol", "lambda_quantile")
+            if k in params
+        }
+        training_overrides = {
+            k: params[k] for k in ("batch_size",) if k in params
+        }
+
+        new_model = replace(cfg.model, **model_overrides) if model_overrides else cfg.model
+        new_asro = replace(cfg.asro, **asro_overrides) if asro_overrides else cfg.asro
+        new_training = replace(cfg.training, **training_overrides) if training_overrides else cfg.training
+
+        logger.info(
+            "Loaded Optuna best params (trial #%d, val_loss=%.4f): %s",
+            data.get("best_trial", -1),
+            data.get("best_value", float("nan")),
+            params,
+        )
+        return replace(cfg, model=new_model, asro=new_asro, training=new_training)
+
+    except Exception as exc:
+        logger.warning("Could not apply Optuna results: %s", exc)
+        return cfg
+
+
 def _persist_tft_metadata(symbol: str, result: dict) -> None:
     """Save TFT model metadata to DB."""
     try: