update ssl and wsl

calculate mean dice for foreground classes
allow loading pre-trained models

Author: taigw

pymic/__init__.py

@@ -1,2 +1,2 @@
 from __future__ import absolute_import
-from . import *
+__version__ = "0.3.1"

pymic/io/__init__.py

@@ -1,2 +1,4 @@
 from __future__ import absolute_import
-from . import *
+from pymic.io.image_read_write import *
+from pymic.io.nifty_dataset import *
+from pymic.io.h5_dataset import *

pymic/net/net2d/unet2d.py

@@ -131,7 +131,6 @@ class Decoder(nn.Module):
     :param class_num: (int) The class number for segmentation task. 
     :param bilinear: (bool) Using bilinear for up-sampling or not. 
         If False, deconvolution will be used for up-sampling.
-    :param multiscale_pred: (bool) Get multi-scale prediction.
     """
     def __init__(self, params):
         super(Decoder, self).__init__()
@@ -140,8 +139,7 @@ def __init__(self, params):
         self.ft_chns   = self.params['feature_chns']
         self.dropout   = self.params['dropout']
         self.n_class   = self.params['class_num']
-        self.bilinear  = self.params.get('bilinear', True)
-        self.mul_pred  = self.params.get('multiscale_pred', False)
+        self.bilinear  = self.params['bilinear']
 
         assert(len(self.ft_chns) == 5 or len(self.ft_chns) == 4)
 
@@ -151,10 +149,6 @@ def __init__(self, params):
         self.up3 = UpBlock(self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], self.dropout[1], self.bilinear) 
         self.up4 = UpBlock(self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], self.dropout[0], self.bilinear) 
         self.out_conv = nn.Conv2d(self.ft_chns[0], self.n_class, kernel_size = 1)
-        if(self.mul_pred):
-            self.out_conv1 = nn.Conv2d(self.ft_chns[1], self.n_class, kernel_size = 1)
-            self.out_conv2 = nn.Conv2d(self.ft_chns[2], self.n_class, kernel_size = 1)
-            self.out_conv3 = nn.Conv2d(self.ft_chns[3], self.n_class, kernel_size = 1)
 
     def forward(self, x):
         if(len(self.ft_chns) == 5):
@@ -169,11 +163,6 @@ def forward(self, x):
         x_d1 = self.up3(x_d2, x1)
         x_d0 = self.up4(x_d1, x0)
         output = self.out_conv(x_d0)
-        if(self.mul_pred):
-            output1 = self.out_conv1(x_d1)
-            output2 = self.out_conv2(x_d2)
-            output3 = self.out_conv3(x_d3)
-            output = [output, output1, output2, output3]
         return output
 
 class UNet2D(nn.Module):

pymic/net_run/agent_abstract.py

@@ -276,7 +276,7 @@ def worker_init_fn(worker_id):
             self.test_loader = torch.utils.data.DataLoader(self.test_set, 
                 batch_size = bn_test, shuffle=False, num_workers= bn_test)
 
-    def create_optimizer(self, params):
+    def create_optimizer(self, params, checkpoint = None):
         """
         Create optimizer based on configuration. 
 
@@ -288,9 +288,9 @@ def create_optimizer(self, params):
             self.optimizer = get_optimizer(opt_params['optimizer'],
                     params, opt_params)
         last_iter = -1
-        if(self.checkpoint is not None):
-            self.optimizer.load_state_dict(self.checkpoint['optimizer_state_dict'])
-            last_iter = self.checkpoint['iteration'] - 1
+        if(checkpoint is not None):
+            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+            last_iter = checkpoint['iteration'] - 1
         if(self.scheduler is None):
             opt_params["last_iter"] = last_iter
             self.scheduler = get_lr_scheduler(self.optimizer, opt_params)

pymic/net_run/agent_seg.py

@@ -29,7 +29,7 @@
 from pymic.transform.trans_dict import TransformDict
 from pymic.util.post_process import PostProcessDict
 from pymic.util.image_process import convert_label
-from pymic.util.general import mixup
+from pymic.util.general import mixup, tensor_shape_match
 
 class SegmentationAgent(NetRunAgent):
     def __init__(self, config, stage = 'train'):
@@ -259,7 +259,8 @@ def train_valid(self):
         ckpt_prefix = self.config['training'].get('ckpt_prefix', None)
         if(ckpt_prefix is None):
             ckpt_prefix = ckpt_dir.split('/')[-1]
-        iter_start  = self.config['training']['iter_start']
+        # iter_start  = self.config['training']['iter_start']       
+        iter_start  = 0     
         iter_max    = self.config['training']['iter_max']
         iter_valid  = self.config['training']['iter_valid']
         iter_save   = self.config['training'].get('iter_save', None)
@@ -274,21 +275,32 @@ def train_valid(self):
         self.max_val_dice = 0.0
         self.max_val_it   = 0
         self.best_model_wts = None 
-        self.checkpoint = None
-        if(iter_start > 0):
-            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)
-            if(len(device_ids) > 1):
-                self.net.module.load_state_dict(self.checkpoint['model_state_dict'])
+        checkpoint = None
+        # initialize the network with pre-trained weights
+        ckpt_init_name = self.config['training'].get('ckpt_init_name', None)
+        ckpt_init_mode = self.config['training'].get('ckpt_init_mode', 0)
+        ckpt_for_optm  = None 
+        if(ckpt_init_name is not None):
+            checkpoint = torch.load(ckpt_dir + "/" + ckpt_init_name, map_location = self.device)
+            pretrained_dict = checkpoint['model_state_dict']
+            model_dict = self.net.module.state_dict() if (len(device_ids) > 1) else self.net.state_dict()
+            pretrained_dict = {k: v for k, v in pretrained_dict.items() if \
+                k in model_dict and tensor_shape_match(pretrained_dict[k], model_dict[k])}
+            logging.info("Initializing the following parameters with pre-trained model")
+            for k in pretrained_dict:
+                logging.info(k)
+            if (len(device_ids) > 1):
+                self.net.module.load_state_dict(pretrained_dict, strict = False)
             else:
-                self.net.load_state_dict(self.checkpoint['model_state_dict'])
-            self.max_val_dice = self.checkpoint.get('valid_pred', 0)
-            # self.max_val_it   = self.checkpoint['iteration']
-            self.max_val_it   = iter_start
-            self.best_model_wts = self.checkpoint['model_state_dict']
-            
-        self.create_optimizer(self.get_parameters_to_update())
+                self.net.load_state_dict(pretrained_dict, strict = False)
+
+            if(ckpt_init_mode > 0): # Load  other information
+                self.max_val_dice = checkpoint.get('valid_pred', 0)
+                iter_start = checkpoint['iteration'] - 1
+                self.max_val_it = iter_start
+                self.best_model_wts = checkpoint['model_state_dict']
+                ckpt_for_optm = checkpoint
+        self.create_optimizer(self.get_parameters_to_update(), ckpt_for_optm)
         self.create_loss_calculator()
 
         self.trainIter  = iter(self.train_loader)

pymic/net_run_nll/nll_co_teaching.py

@@ -32,7 +32,7 @@ def forward(self, x):
         if(self.training):
           return out1, out2
         else:
-          return (out1 + out2) / 3
+          return (out1 + out2) / 2
 
 class NLLCoTeaching(SegmentationAgent):
     """
@@ -144,13 +144,13 @@ def training(self):
         train_avg_loss1 = train_loss1 / iter_valid
         train_avg_loss2 = train_loss2 / iter_valid
         train_cls_dice = np.asarray(train_dice_list).mean(axis = 0)
-        train_avg_dice = train_cls_dice.mean()
+        train_avg_dice = train_cls_dice[1:].mean()
 
         train_scalers = {'loss': (train_avg_loss1 + train_avg_loss2) / 2, 
             'loss1':train_avg_loss1, 'loss2': train_avg_loss2,
             'loss_no_select1':train_avg_loss_no_select1, 
             'loss_no_select2':train_avg_loss_no_select2,
-            'select_ratio':remb_ratio, 'avg_dice':train_avg_dice, 'class_dice': train_cls_dice}
+            'select_ratio':remb_ratio, 'avg_fg_dice':train_avg_dice, 'class_dice': train_cls_dice}
         return train_scalers
 
     def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
@@ -159,7 +159,7 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
         loss_no_select_scalar  = {'net1':train_scalars['loss_no_select1'],
                                   'net2':train_scalars['loss_no_select2']}
 
-        dice_scalar ={'train':train_scalars['avg_dice'], 'valid':valid_scalars['avg_dice']}
+        dice_scalar ={'train':train_scalars['avg_fg_dice'], 'valid':valid_scalars['avg_fg_dice']}
         self.summ_writer.add_scalars('loss', loss_scalar, glob_it)
         self.summ_writer.add_scalars('loss_no_select', loss_no_select_scalar, glob_it)
         self.summ_writer.add_scalars('select_ratio', {'select_ratio':train_scalars['select_ratio']}, glob_it)
@@ -171,9 +171,9 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
                 'valid':valid_scalars['class_dice'][c]}
             self.summ_writer.add_scalars('class_{0:}_dice'.format(c), cls_dice_scalar, glob_it)
 
-        logging.info('train loss {0:.4f}, avg dice {1:.4f} '.format(
-            train_scalars['loss'], train_scalars['avg_dice']) + "[" + \
+        logging.info('train loss {0:.4f}, avg foreground dice {1:.4f} '.format(
+            train_scalars['loss'], train_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in train_scalars['class_dice']) + "]")        
-        logging.info('valid loss {0:.4f}, avg dice {1:.4f} '.format(
-            valid_scalars['loss'], valid_scalars['avg_dice']) + "[" + \
+        logging.info('valid loss {0:.4f}, avg foreground dice {1:.4f} '.format(
+            valid_scalars['loss'], valid_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in valid_scalars['class_dice']) + "]") 

pymic/net_run_nll/nll_dast.py

@@ -5,7 +5,6 @@
 import numpy as np 
 import torch.nn as nn
 import torchvision.transforms as transforms
-from torch.optim import lr_scheduler
 from pymic.io.nifty_dataset import NiftyDataset
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
@@ -257,11 +256,11 @@ def training(self):
         train_avg_loss_sup = train_loss_sup / iter_valid
         train_avg_loss_reg = train_loss_reg / iter_valid
         train_cls_dice = np.asarray(train_dice_list).mean(axis = 0)
-        train_avg_dice = train_cls_dice.mean()
+        train_avg_dice = train_cls_dice[1:].mean()
 
         train_scalers = {'loss': train_avg_loss, 'loss_sup':train_avg_loss_sup,
             'loss_reg':train_avg_loss_reg, 'regular_w':w_dbc,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
 
     def train_valid(self):

pymic/net_run_nll/nll_main.py

@@ -28,8 +28,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)
     nll_method = config['noisy_label_learning']['nll_method']

pymic/net_run_nll/nll_trinet.py

@@ -140,13 +140,13 @@ def training(self):
         train_avg_loss1 = train_loss1 / iter_valid
         train_avg_loss2 = train_loss2 / iter_valid
         train_cls_dice = np.asarray(train_dice_list).mean(axis = 0)
-        train_avg_dice = train_cls_dice.mean()
+        train_avg_dice = train_cls_dice[1:].mean()
 
         train_scalers = {'loss': (train_avg_loss1 + train_avg_loss2) / 2, 
             'loss1':train_avg_loss1, 'loss2': train_avg_loss2,
             'loss_no_select1':train_avg_loss_no_select1, 
             'loss_no_select2':train_avg_loss_no_select2,
-            'select_ratio':remb_ratio, 'avg_dice':train_avg_dice, 'class_dice': train_cls_dice}
+            'select_ratio':remb_ratio, 'avg_fg_dice':train_avg_dice, 'class_dice': train_cls_dice}
         return train_scalers
 
     def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
@@ -155,7 +155,7 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
         loss_no_select_scalar  = {'net1':train_scalars['loss_no_select1'],
                                   'net2':train_scalars['loss_no_select2']}
 
-        dice_scalar ={'train':train_scalars['avg_dice'], 'valid':valid_scalars['avg_dice']}
+        dice_scalar ={'train':train_scalars['avg_fg_dice'], 'valid':valid_scalars['avg_fg_dice']}
         self.summ_writer.add_scalars('loss', loss_scalar, glob_it)
         self.summ_writer.add_scalars('loss_no_select', loss_no_select_scalar, glob_it)
         self.summ_writer.add_scalars('select_ratio', {'select_ratio':train_scalars['select_ratio']}, glob_it)
@@ -167,9 +167,9 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
                 'valid':valid_scalars['class_dice'][c]}
             self.summ_writer.add_scalars('class_{0:}_dice'.format(c), cls_dice_scalar, glob_it)
 
-        logging.info('train loss {0:.4f}, avg dice {1:.4f} '.format(
-            train_scalars['loss'], train_scalars['avg_dice']) + "[" + \
+        logging.info('train loss {0:.4f}, avg foregournd dice {1:.4f} '.format(
+            train_scalars['loss'], train_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in train_scalars['class_dice']) + "]")        
-        logging.info('valid loss {0:.4f}, avg dice {1:.4f} '.format(
-            valid_scalars['loss'], valid_scalars['avg_dice']) + "[" + \
+        logging.info('valid loss {0:.4f}, avg foreground dice {1:.4f} '.format(
+            valid_scalars['loss'], valid_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in valid_scalars['class_dice']) + "]") 

pymic/net_run_ssl/ssl_cps.py

@@ -166,8 +166,8 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
             self.summ_writer.add_scalars('class_{0:}_dice'.format(c), cls_dice_scalar, glob_it)
 
         logging.info('train loss {0:.4f}, avg dice {1:.4f} '.format(
-            train_scalars['loss'], train_scalars['avg_dice']) + "[" + \
+            train_scalars['loss'], train_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in train_scalars['class_dice']) + "]")        
         logging.info('valid loss {0:.4f}, avg dice {1:.4f} '.format(
-            valid_scalars['loss'], valid_scalars['avg_dice']) + "[" + \
+            valid_scalars['loss'], valid_scalars['avg_fg_dice']) + "[" + \
             ' '.join("{0:.4f}".format(x) for x in valid_scalars['class_dice']) + "]")