Merge branch 'REPST' of https://github.zhang-heng.com/czzhangheng/TrafficWheel into REPST

config/REPST/PEMS-BAY.yaml

@@ -43,13 +43,14 @@ train:
   early_stop_patience: 15
   epochs: 100
   grad_norm: false
-  loss_func: mae
+  loss_func: mse
+  scheduler: cosineAnnealingLR
+  tmax: 5
   lr_decay: true
   lr_decay_rate: 0.3
   lr_decay_step: "5,20,40,70"
   lr_init: 0.003
   max_grad_norm: 5
-  real_value: true
   weight_decay: 0
   debug: false
   output_dim: 1

utils/initializer.py

@@ -23,6 +23,9 @@ def init_model(args):
 
 
 def init_optimizer(model, args):
+    optim = args.get("optimizer", "Adam")
+    match optim :
+        case "Adam":
             optimizer = torch.optim.Adam(
             params=model.parameters(),
             lr=args["lr_init"],
@@ -31,13 +34,18 @@ def init_optimizer(model, args):
             amsgrad=False,
         )
     
+    scheduler = args.get("scheduler", "multistepLR")
+    match scheduler:
+        case "multistepLR":
             lr_scheduler = None
             if args["lr_decay"]:
                 lr_decay_steps = [int(step) for step in args["lr_decay_step"].split(",")]
                 lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
                     optimizer=optimizer, milestones=lr_decay_steps, gamma=args["lr_decay_rate"]
                 )
-
+        case "cosineAnnealingLR":
+            T_max = args.get("tmax", 5)
+            lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=T_max, eta_min=1e-8)
     return optimizer, lr_scheduler