update learning rate scheduler

to keep consistent with pytorch, one step in  learning rate scheduler corresponds to an epoch on the training set

Author: taigw

pymic/net_run/agent_seg.py

@@ -162,10 +162,6 @@ def training(self):
             loss = self.get_loss_value(data, outputs, labels_prob)
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
-
             train_loss = train_loss + loss.item()
             # get dice evaluation for each class
             if(isinstance(outputs, tuple) or isinstance(outputs, list)):
@@ -219,10 +215,6 @@ def validation(self):
         valid_avg_loss = np.asarray(valid_loss_list).mean()
         valid_cls_dice = np.asarray(valid_dice_list).mean(axis = 0)
         valid_avg_dice = valid_cls_dice.mean()
-        
-        if(isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-            self.scheduler.step(valid_avg_dice)
-
         valid_scalers = {'loss': valid_avg_loss, 'avg_dice': valid_avg_dice,\
             'class_dice': valid_cls_dice}
         return valid_scalers
@@ -300,9 +292,13 @@ def train_valid(self):
             t0 = time.time()
             train_scalars = self.training()
             t1 = time.time()
-            
             valid_scalars = self.validation()
             t2 = time.time()
+            if(isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
+                self.scheduler.step(valid_scalars['avg_dice'])
+            else:
+                self.scheduler.step()
+
             self.glob_it = it + iter_valid
             logging.info("\n{0:} it {1:}".format(str(datetime.now())[:-7], self.glob_it))
             logging.info('learning rate {0:}'.format(lr_value))

pymic/net_run/get_optimizer.py

@@ -38,20 +38,28 @@ def get_optimizer(name, net_params, optim_params):
 
 def get_lr_scheduler(optimizer, sched_params):
     name = sched_params["lr_scheduler"]
+    val_it = sched_params["iter_valid"]
     if(name is None):
         return None
-    lr_gamma = sched_params["lr_gamma"]
     if(keyword_match(name, "ReduceLROnPlateau")):
         patience_it = sched_params["ReduceLROnPlateau_patience".lower()]
-        val_it = sched_params["iter_valid"]
-        patience = patience_it / val_it
+        patience = patience_it / val_it   
+        lr_gamma = sched_params["lr_gamma"]
         scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
             mode = "max", factor=lr_gamma, patience = patience)
     elif(keyword_match(name, "MultiStepLR")):
         lr_milestones = sched_params["lr_milestones"]
-        last_iter     = sched_params["last_iter"]
+        lr_milestones = [int(item / val_it) for item in lr_milestones]
+        epoch_last    = sched_params["last_iter"] / val_it
+        lr_gamma  = sched_params["lr_gamma"]
         scheduler = lr_scheduler.MultiStepLR(optimizer,
-                    lr_milestones, lr_gamma, last_iter)
+                    lr_milestones, lr_gamma, epoch_last)
+    elif(keyword_match(name, "CosineAnnealingLR")):
+        epoch_max  = sched_params["iter_max"] / val_it
+        epoch_last = sched_params["last_iter"] / val_it
+        lr_min = sched_params.get("lr_min", 0)
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
+                    epoch_max, lr_min, epoch_last)
     else:
         raise ValueError("unsupported lr scheduler {0:}".format(name))
     return scheduler

pymic/net_run_nll/nll_co_teaching.py

@@ -128,9 +128,6 @@ def training(self):
 
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
 
             train_loss_no_select1 = train_loss_no_select1 + loss1.mean().item()
             train_loss_no_select2 = train_loss_no_select2 + loss2.mean().item()

pymic/net_run_nll/nll_dast.py

@@ -239,9 +239,6 @@ def training(self):
 
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
 
             train_loss = train_loss + loss.item()
             train_loss_sup = train_loss_sup + loss_sup.item()

pymic/net_run_nll/nll_trinet.py

@@ -8,7 +8,6 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.optim import lr_scheduler
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
@@ -125,9 +124,6 @@ def training(self):
 
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
 
             train_loss_no_select1 = train_loss_no_select1 + loss1.mean().item()
             train_loss_no_select2 = train_loss_no_select2 + loss2.mean().item()

pymic/net_run_ssl/ssl_cct.py

@@ -5,7 +5,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-from torch.optim import lr_scheduler
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
@@ -139,9 +138,7 @@ def training(self):
 
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
+
             train_loss = train_loss + loss.item()
             train_loss_sup = train_loss_sup + loss_sup.item()
             train_loss_reg = train_loss_reg + loss_reg.item() 

pymic/net_run_ssl/ssl_cps.py

@@ -4,7 +4,6 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from torch.optim import lr_scheduler
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
@@ -26,7 +25,7 @@ def forward(self, x):
         if(self.training):
           return out1, out2
         else:
-          return (out1 + out2) / 3
+          return (out1 + out2) / 2
 
 class SSLCPS(SSLSegAgent):
     """
@@ -117,9 +116,6 @@ def training(self):
 
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
 
             train_loss = train_loss + loss.item()
             train_loss_sup1  = train_loss_sup1 + loss_sup1.item()

pymic/net_run_ssl/ssl_em.py

@@ -3,7 +3,6 @@
 import logging
 import numpy as np
 import torch
-from torch.optim import lr_scheduler
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
@@ -83,9 +82,6 @@ def training(self):
             # if (self.config['training']['use'])
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
 
             train_loss = train_loss + loss.item()
             train_loss_sup = train_loss_sup + loss_sup.item()

pymic/net_run_ssl/ssl_mt.py

@@ -3,7 +3,6 @@
 import logging
 import torch
 import numpy as np
-from torch.optim import lr_scheduler
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
@@ -102,9 +101,6 @@ def training(self):
 
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
 
             # update EMA
             alpha = ssl_cfg.get('ema_decay', 0.99)

pymic/net_run_ssl/ssl_uamt.py

@@ -3,7 +3,6 @@
 import logging
 import torch
 import numpy as np
-from torch.optim import lr_scheduler
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
@@ -104,10 +103,6 @@ def training(self):
 
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
-
 
             # update EMA
             alpha = ssl_cfg.get('ema_decay', 0.99)