fix learning rate bug

cybershiptrooper · Sep 4, 2024 · 84fdd6f · 84fdd6f
1 parent 73c1663
commit 84fdd6f
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/iit/model_pairs/base_model_pair.py b/iit/model_pairs/base_model_pair.py
@@ -247,7 +247,11 @@ def train(
 
         early_stop = training_args["early_stop"]
 
-        optimizer = training_args['optimizer_cls'](self.ll_model.parameters(), **training_args['optimizer_kwargs'])
+        optimizer = training_args['optimizer_cls'](
+            self.ll_model.parameters(), 
+            lr=training_args["lr"], 
+            **training_args['optimizer_kwargs']
+        )
         loss_fn = self.loss_fn
         scheduler_cls = training_args.get("lr_scheduler", None)
         scheduler_kwargs = training_args.get("scheduler_kwargs", {})

diff --git a/iit/model_pairs/iit_behavior_model_pair.py b/iit/model_pairs/iit_behavior_model_pair.py
@@ -21,7 +21,7 @@ def __init__(
     ):
         default_training_args = {
             "atol": 5e-2,
-            "use_single_loss": False,
+            "use_single_loss": True,
             "iit_weight": 1.0,
             "behavior_weight": 1.0,
             "val_IIA_sampling": "random",  # random or all

diff --git a/iit/model_pairs/iit_model_pair.py b/iit/model_pairs/iit_model_pair.py
@@ -34,10 +34,11 @@ def __init__(
             "scheduler_kwargs": {},
             "clip_grad_norm": 1.0,
             "seed": 0,
+            "lr": 0.001,
             "detach_while_caching": True,
             "optimizer_cls": t.optim.Adam,
             "optimizer_kwargs" : {
-                "lr": 0.001,
+                # "betas": (0.9, 0.9)
             },
         }
         training_args = {**default_training_args, **training_args}