optimizer name string as command line input

marrlab · Nov 29, 2024 · 495c314 · 495c314
1 parent 6da7252
commit 495c314
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/domainlab/algos/trainers/a_trainer.py b/domainlab/algos/trainers/a_trainer.py
@@ -14,7 +14,8 @@ def mk_opt(model, aconf):
     create optimizer
     """
     if model._decoratee is None:
-        optimizer = optim.Adam(model.parameters(), lr=aconf.lr)
+        class_opt = getattr(optim, aconf.opt)
+        optimizer = class_opt(model.parameters(), lr=aconf.lr)
     else:
         var1 = model.parameters()
         var2 = model._decoratee.parameters()

diff --git a/domainlab/arg_parser.py b/domainlab/arg_parser.py
@@ -257,6 +257,13 @@ def mk_parser_main():
         "Default is zoutput/benchmarks/shell_benchmark",
     )
 
+    parser.add_argument(
+        "--opt",
+        type=str,
+        default="Adam",
+        help="name of pytorch optimizer",
+    )
+
     parser.add_argument(
         "--param_idx",
         type=bool,