bugfix(sup): fix wrong lr times

ZombaSY · May 6, 2022 · 07415ac · 07415ac
1 parent 08a6d48
commit 07415ac
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train_sup.py b/train_sup.py
@@ -97,6 +97,7 @@ def main():
     # Optimizer and lr decay scheduler
     cfg_trainer = cfg["trainer"]
     cfg_optim = cfg_trainer["optimizer"]
+    times = 10 if "pascal" in cfg["dataset"]["type"] else 1
 
     params_list = []
     for module in modules_back:
@@ -105,7 +106,7 @@ def main():
         )
     for module in modules_head:
         params_list.append(
-            dict(params=module.parameters(), lr=cfg_optim["kwargs"]["lr"] * 10)
+            dict(params=module.parameters(), lr=cfg_optim["kwargs"]["lr"] * times)
         )
 
     optimizer = get_optimizer(params_list, cfg_optim)