add adamW optimizer

Author: KMarshallX

config/adapt_config.py

@@ -36,7 +36,7 @@
 adapt_parser.add_argument('--patch_size', nargs=3, type=int, default=(64, 64, 64),
                         help='Expected size for training (x y z)')
 # optimizer type, available: [sgd, adam]
-adapt_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam]')
+adapt_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam, adamw]')
 # loss metric type, available: [bce, dice, tver]
 adapt_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 

config/angiboost_config.py

@@ -31,7 +31,7 @@
 # expected size of the training patch
 angiboost_parser.add_argument('--osz', nargs=3, type=int, default=(64, 64, 64),
                         help='Expected size of the training patch (x y z)')
-angiboost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam]")
+angiboost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam, adamw]")
 angiboost_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 
 # Optimizer tuning

config/boost_config.py

@@ -30,7 +30,7 @@
 # expected size of the training patch
 boost_parser.add_argument('--patch_size', nargs=3, type=int, default=(64, 64, 64),
                         help='Expected size of the training patch (x y z)')
-boost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam]")
+boost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam, adamw]")
 boost_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 
 # Optimizer tuning

config/train_config.py

@@ -29,7 +29,7 @@
 # expected patch size for training
 train_parser.add_argument('--output_size', nargs=3, type=int, default=(64, 64, 64), help='Expected patch size for training (x y z)')
 # optimizer type, available: [sgd, adam]
-train_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam]')
+train_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam, adamw]')
 # loss metric type, available: [bce, dice, tver]
 train_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 

library/loss_func.py

@@ -150,7 +150,8 @@ def choose_optimizer(optim_name: str, model_params: Any, lr: float) -> Any:
     """
     optimizer_registry = {
         'sgd': lambda: torch.optim.SGD(model_params, lr),
-        'adam': lambda: torch.optim.Adam(model_params, lr)
+        'adam': lambda: torch.optim.Adam(model_params, lr),
+        'adamw': lambda: torch.optim.AdamW(model_params, lr)
     }
 
     if optim_name not in optimizer_registry: