update/put optimization defaults in one place (#288)

Author: themattinthehatt

lightning_pose/models/base.py

@@ -33,7 +33,7 @@
 
 DEFAULT_LR_SCHEDULER_PARAMS = OmegaConf.create(
     {
-        "milestones": [100, 200, 300],
+        "milestones": [150, 200, 250],
         "gamma": 0.5,
     }
 )

lightning_pose/utils/scripts.py

@@ -40,6 +40,10 @@
     SemiSupervisedHeatmapTrackerMHCRNN,
     SemiSupervisedRegressionTracker,
 )
+from lightning_pose.models.base import (
+    _apply_defaults_for_lr_scheduler_params,
+    _apply_defaults_for_optimizer_params,
+)
 from lightning_pose.utils import io as io_utils
 from lightning_pose.utils.pca import KeypointPCA
 
@@ -345,15 +349,15 @@ def get_model(
     """Create model: regression or heatmap based, supervised or semi-supervised."""
 
     optimizer = cfg.training.get("optimizer", "Adam")
-    optimizer_params = cfg.training.get(
-        "optimizer_params",
-        OmegaConf.create({"learning_rate": 1e-4}),
+    optimizer_params = _apply_defaults_for_optimizer_params(
+        optimizer,
+        cfg.training.get("optimizer_params"),
     )
 
-    lr_scheduler = cfg.training.lr_scheduler
-
-    lr_scheduler_params = OmegaConf.to_object(
-        cfg.training.lr_scheduler_params[lr_scheduler]
+    lr_scheduler = cfg.training.get("lr_scheduler", "multisteplr")
+    lr_scheduler_params = _apply_defaults_for_lr_scheduler_params(
+        lr_scheduler,
+        cfg.training.get("lr_scheduler_params", {}).get(f"{lr_scheduler}")
     )
 
     semi_supervised = io_utils.check_if_semi_supervised(cfg.model.losses_to_use)