update main files

lr_schduler update after each validation for classification task;
update main files for logging bug
update init
fix bug for alpha in mean teacher

Author: taigw

README.md

@@ -15,9 +15,9 @@ BibTeX entry:
     author = {Guotai Wang and Xiangde Luo and Ran Gu and Shuojue Yang and Yijie Qu and Shuwei Zhai and Qianfei Zhao and Kang Li and Shaoting Zhang},
     title = {{PyMIC: A deep learning toolkit for annotation-efficient medical image segmentation}},
     year = {2023},
-    url = {http://arxiv.org/abs/2208.09350},
+    url = {https://doi.org/10.1016/j.cmpb.2023.107398},
     journal = {Computer Methods and Programs in Biomedicine},
-    volume = {February},
+    volume = {231},
     pages = {107398},
     }
 

pymic/net_run/agent_cls.py

@@ -157,9 +157,6 @@ def training(self):
             loss = self.get_loss_value(data, outputs, labels)
             loss.backward()
             self.optimizer.step()
-            if(self.scheduler is not None and \
-                not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-                self.scheduler.step()
 
             # statistics
             sample_num   += labels.size(0)
@@ -183,7 +180,7 @@ def validation(self):
                 inputs = self.convert_tensor_type(data['image'])
                 labels = self.convert_tensor_type(data['label_prob'])            
                 inputs, labels = inputs.to(self.device), labels.to(self.device)
-                self.optimizer.zero_grad()
+                # self.optimizer.zero_grad()
                 # forward + backward + optimize
                 outputs = self.net(inputs)
                 loss = self.get_loss_value(data, outputs, labels)
@@ -196,20 +193,17 @@ def validation(self):
         avg_loss = running_loss / sample_num
         avg_score= running_score.double() / sample_num
         metrics  = self.config['training'].get("evaluation_metric", "accuracy")
-        if(isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau)):
-            self.scheduler.step(avg_score)
         valid_scalers = {'loss': avg_loss, metrics: avg_score}
         return valid_scalers
 
     def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
-        metrics =self.config['training'].get("evaluation_metric", "accuracy")
+        metrics = self.config['training'].get("evaluation_metric", "accuracy")
         loss_scalar ={'train':train_scalars['loss'], 'valid':valid_scalars['loss']}
         acc_scalar  ={'train':train_scalars[metrics],'valid':valid_scalars[metrics]}
         self.summ_writer.add_scalars('loss', loss_scalar, glob_it)
         self.summ_writer.add_scalars(metrics, acc_scalar, glob_it)
         self.summ_writer.add_scalars('lr', {"lr": lr_value}, glob_it)
 
-        logging.info("{0:} it {1:}".format(str(datetime.now())[:-7], glob_it))
         logging.info('train loss {0:.4f}, avg {1:} {2:.4f}'.format(
             train_scalars['loss'], metrics, train_scalars[metrics]))
         logging.info('valid loss {0:.4f}, avg {1:} {2:.4f}'.format(
@@ -251,7 +245,10 @@ def train_valid(self):
             checkpoint_file = "{0:}/{1:}_{2:}.pt".format(ckpt_dir, ckpt_prefix, iter_start)
             self.checkpoint = torch.load(checkpoint_file, map_location = self.device)
             assert(self.checkpoint['iteration'] == iter_start)
-            self.net.load_state_dict(self.checkpoint['model_state_dict'])
+            if(len(device_ids) > 1):
+                self.net.module.load_state_dict(self.checkpoint['model_state_dict'])
+            else:
+                self.net.load_state_dict(self.checkpoint['model_state_dict'])
             self.max_val_score  = self.checkpoint.get('valid_pred', 0)
             self.max_val_it     = self.checkpoint['iteration']
             self.best_model_wts = self.checkpoint['model_state_dict']
@@ -266,15 +263,28 @@ def train_valid(self):
         self.glob_it = iter_start
         for it in range(iter_start, iter_max, iter_valid):
             lr_value = self.optimizer.param_groups[0]['lr']
+            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[metrics])
+            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))
+            logging.info("training/validation time: {0:.2f}s/{1:.2f}s".format(t1-t0, t2-t1))
             self.write_scalars(train_scalars, valid_scalars, lr_value, self.glob_it)
-
             if(valid_scalars[metrics] > self.max_val_score):
                 self.max_val_score = valid_scalars[metrics]
                 self.max_val_it    = self.glob_it
-                self.best_model_wts = copy.deepcopy(self.net.state_dict())
+                if(len(device_ids) > 1):
+                    self.best_model_wts = copy.deepcopy(self.net.module.state_dict())
+                else:
+                    self.best_model_wts = copy.deepcopy(self.net.state_dict())
 
             stop_now = True if(early_stop_it is not None and \
                 self.glob_it - self.max_val_it > early_stop_it) else False
@@ -306,7 +316,6 @@ def train_valid(self):
             self.max_val_it, metrics, self.max_val_score))
         self.summ_writer.close()
 
-
     def infer(self):
         device_ids = self.config['testing']['gpus']
         device = torch.device("cuda:{0:}".format(device_ids[0]))

pymic/net_run/agent_seg.py

@@ -96,17 +96,23 @@ def create_loss_calculator(self):
             raise ValueError("Undefined loss function {0:}".format(loss_name))
         else:
             base_loss = self.loss_dict[loss_name](self.config['training'])
-        if(self.config['network'].get('deep_supervise', False)):
-            weight = self.config['network'].get('deep_supervise_weight', None)
-            params = {'deep_supervise_weight': weight, 'base_loss':base_loss}
+        if(self.config['training'].get('deep_supervise', False)):
+            weight = self.config['training'].get('deep_supervise_weight', None)
+            mode   = self.config['training'].get('deep_supervise_mode', 2)
+            params = {'deep_supervise_weight': weight, 
+                      'deep_supervise_mode': mode, 
+                      'base_loss':base_loss}
             self.loss_calculator = DeepSuperviseLoss(params)
         else:
             self.loss_calculator = base_loss
 
     def get_loss_value(self, data, pred, gt, param = None):
         loss_input_dict = {'prediction':pred, 'ground_truth': gt}
         if data.get('pixel_weight', None) is not None:
-            loss_input_dict['pixel_weight'] = data['pixel_weight'].to(pred.device)
+            if(isinstance(pred, tuple) or isinstance(pred, list)):
+                loss_input_dict['pixel_weight'] = data['pixel_weight'].to(pred[0].device)
+            else:
+                loss_input_dict['pixel_weight'] = data['pixel_weight'].to(pred.device)
         loss_value = self.loss_calculator(loss_input_dict)
         return loss_value
 
@@ -122,7 +128,7 @@ def set_postprocessor(self, postprocessor):
     def training(self):
         class_num   = self.config['network']['class_num']
         iter_valid  = self.config['training']['iter_valid']
-        mixup_prob  = self.config['training'].get('mixup_probability', 0.5)
+        mixup_prob  = self.config['training'].get('mixup_probability', 0.0)
         train_loss  = 0
         train_dice_list = []
         self.net.train()
@@ -135,7 +141,7 @@ def training(self):
             # get the inputs
             inputs      = self.convert_tensor_type(data['image'])
             labels_prob = self.convert_tensor_type(data['label_prob'])                 
-            if(random() < mixup_prob):
+            if(mixup_prob > 0 and random() < mixup_prob):
                 inputs, labels_prob = mixup(inputs, labels_prob) 
 
             # # for debug
@@ -246,7 +252,10 @@ def train_valid(self):
         else:
             self.device = torch.device("cuda:{0:}".format(device_ids[0]))
         self.net.to(self.device)
+        
         ckpt_dir    = self.config['training']['ckpt_save_dir']
+        if(ckpt_dir[-1] == "/"):
+            ckpt_dir = ckpt_dir[:-1]
         ckpt_prefix = self.config['training'].get('ckpt_prefix', None)
         if(ckpt_prefix is None):
             ckpt_prefix = ckpt_dir.split('/')[-1]

pymic/net_run/get_optimizer.py

@@ -59,7 +59,7 @@ def get_lr_scheduler(optimizer, sched_params):
         scheduler = lr_scheduler.MultiStepLR(optimizer,
                     lr_milestones, lr_gamma, epoch_last)
     elif(keyword_match(name, "StepLR")):
-        lr_step = sched_params["lr_step"] / val_it
+        lr_step   = sched_params["lr_step"] / val_it
         lr_gamma  = sched_params["lr_gamma"]
         scheduler = lr_scheduler.StepLR(optimizer,
                     lr_step, lr_gamma, epoch_last)

pymic/net_run_ssl/__init__.py

@@ -1,2 +1,8 @@
 from __future__ import absolute_import
-from . import *
+from pymic.net_run_ssl.ssl_abstract import *
+from pymic.net_run_ssl.ssl_cct import *
+from pymic.net_run_ssl.ssl_cps import *
+from pymic.net_run_ssl.ssl_em import *
+from pymic.net_run_ssl.ssl_mt import *
+from pymic.net_run_ssl.ssl_uamt import *
+from pymic.net_run_ssl.ssl_urpc import *

pymic/net_run_ssl/ssl_main.py

@@ -35,8 +35,12 @@ def main():
     log_dir  = config['training']['ckpt_save_dir']
     if(not os.path.exists(log_dir)):
         os.mkdir(log_dir)
-    logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
-                        format='%(message)s')
+    if sys.version.startswith("3.9"):
+        logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
+                            format='%(message)s', force=True) # for python 3.9
+    else:
+        logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
+                            format='%(message)s') # for python 3.6
     logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
     logging_config(config)
     ssl_method = config['semi_supervised_learning']['ssl_method']

pymic/net_run_ssl/ssl_mt.py

@@ -104,7 +104,7 @@ def training(self):
 
             # update EMA
             alpha = ssl_cfg.get('ema_decay', 0.99)
-            alpha = min(1 - 1 / (iter_max + 1), alpha)
+            alpha = min(1 - 1 / (self.glob_it / iter_valid + 1), alpha)
             for ema_param, param in zip(self.net_ema.parameters(), self.net.parameters()):
                 ema_param.data.mul_(alpha).add_(1 - alpha, param.data)
 

pymic/net_run_ssl/ssl_uamt.py

@@ -106,7 +106,7 @@ def training(self):
 
             # update EMA
             alpha = ssl_cfg.get('ema_decay', 0.99)
-            alpha = min(1 - 1 / (iter_max + 1), alpha)
+            alpha = min(1 - 1 / (self.glob_it / iter_valid + 1), alpha)
             for ema_param, param in zip(self.net_ema.parameters(), self.net.parameters()):
                 ema_param.data.mul_(alpha).add_(1 - alpha, param.data)
 

pymic/net_run_wsl/wsl_main.py

@@ -34,8 +34,12 @@ def main():
     log_dir  = config['training']['ckpt_save_dir']
     if(not os.path.exists(log_dir)):
         os.mkdir(log_dir)
-    logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
-                        format='%(message)s')
+    if sys.version.startswith("3.9"):
+        logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
+                            format='%(message)s', force=True) # for python 3.9
+    else:
+        logging.basicConfig(filename=log_dir+"/log_{0:}.txt".format(stage), level=logging.INFO,
+                            format='%(message)s') # for python 3.6
     logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
     logging_config(config)
     wsl_method = config['weakly_supervised_learning']['wsl_method']

pymic/transform/__init__.py

@@ -1,2 +1,13 @@
+# -*- coding: utf-8 -*-
 from __future__ import absolute_import
-from . import *
+from pymic.transform.intensity import  *
+from pymic.transform.flip import *
+from pymic.transform.pad import *
+from pymic.transform.rotate import *
+from pymic.transform.rescale import *
+from pymic.transform.transpose import *
+from pymic.transform.threshold import * 
+from pymic.transform.normalize import *
+from pymic.transform.crop import *
+from pymic.transform.label_convert import * 
+from pymic.transform.trans_dict import TransformDict