feat: weight decay hparams tuning

DimensionSTP · Jul 19, 2024 · 7b48c17 · 7b48c17
1 parent 4fe1a80
commit 7b48c17
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/tuners/rhythm_tuner.py b/src/tuners/rhythm_tuner.py
@@ -109,6 +109,13 @@ def optuna_objective(
                 high=self.hparams.lr.high,
                 log=self.hparams.lr.log,
             )
+        if self.hparams.weight_decay:
+            params["weight_decay"] = trial.suggest_float(
+                name="weight_decay",
+                low=self.hparams.weight_decay.low,
+                high=self.hparams.weight_decay.high,
+                log=self.hparams.weight_decay.log,
+            )
         if self.hparams.period:
             params["period"] = trial.suggest_int(
                 name="period",
@@ -135,6 +142,7 @@ def optuna_objective(
             model=model,
             strategy=self.module_params.strategy,
             lr=params["lr"],
+            weight_decay=params["weight_decay"],
             period=params["period"],
             eta_min=params["eta_min"],
             interval=self.module_params.interval,