Repository: HiLab-git/PyMIC (language: python)

Issue:
Dev



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 3ce5929..ac7f9b6 100644
--- a/README.md
+++ b/README.md
@@ -2,17 +2,28 @@
 
 PyMIC is a pytorch-based toolkit for medical image computing with annotation-efficient deep learning. Despite that pytorch is a fantastic platform for deep learning, using it for medical image computing is not straightforward as medical images are often with high dimension and large volume, multiple modalities and difficulies in annotating. This toolkit is developed to facilitate medical image computing researchers so that training and testing deep learning models become easier. It is very friendly to researchers who are new to this area. Even without writing any code, you can use PyMIC commands to train and test a model by simply editing configuration files.  PyMIC is developed to support learning with imperfect labels, including semi-supervised and weakly supervised learning, and learning with noisy annotations.
 
-Currently PyMIC supports 2D/3D medical image classification and segmentation, and it is still under development. It was originally developed for COVID-19 pneumonia lesion segmentation from CT images. If you use this toolkit, please cite the following paper:
+Currently PyMIC supports 2D/3D medical image classification and segmentation, and it is still under development. If you use this toolkit, please cite the following paper:
 
-*  G. Wang, X. Liu, C. Li, Z. Xu, J. Ruan, H. Zhu, T. Meng, K. Li, N. Huang, S. Zhang. 
-[A Noise-robust Framework for Automatic Segmentation of COVID-19 Pneumonia Lesions from CT Images.][tmi2020] IEEE Transactions on Medical Imaging. 39(8):2653-2663, 2020. DOI: [10.1109/TMI.2020.3000314][tmi2020]
+* G. Wang, X. Luo, R. Gu, S. Yang, Y. Qu, S. Zhai, Q. Zhao, K. Li, S. Zhang. (2023). 
+[PyMIC: A deep learning toolkit for annotation-efficient medical image segmentation.][arxiv2022] Computer Methods and Programs in Biomedicine (CMPB). February 2023, 107398.
 
-[tmi2020]:https://ieeexplore.ieee.org/document/9109297
+[arxiv2022]:http://arxiv.org/abs/2208.09350
 
+BibTeX entry:
+
+    @article{Wang2022pymic,
+    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 = {https://doi.org/10.1016/j.cmpb.2023.107398},
+    journal = {Computer Methods and Programs in Biomedicine},
+    volume = {231},
+    pages = {107398},
+    }
 
 # Features
 PyMIC provides flixible modules for medical image computing tasks including classification and segmentation. It currently provides the following functions:
-* Support for annotation-efficient image segmentation, especially for semi-supervised, weakly-supervised and noisy-label learning.
+* Support for annotation-efficient image segmentation, especially for semi-supervised, self-supervised, weakly-supervised and noisy-label learning.
 * User friendly: For beginners, you only need to edit the configuration files for model training and inference, without writing code. For advanced users, you can customize different modules (networks, loss functions, training pipeline, etc) and easily integrate them into PyMIC.
 * Easy-to-use I/O interface to read and write different 2D and 3D images.
 * Various data pre-processing/transformation methods before sending a tensor into a network.
@@ -36,10 +47,10 @@ Run the following command to install the latest released version of PyMIC:
 ```bash
 pip install PYMIC
 ```
-To install a specific version of PYMIC such as 0.3.0, run:
+To install a specific version of PYMIC such as 0.4.0, run:
 
 ```bash
-pip install PYMIC==0.3.0
+pip install PYMIC==0.4.0
 ```
 Alternatively, you can download the source code for the latest version. Run the following command to compile and install:
 
diff --git a/pymic/__init__.py b/pymic/__init__.py
index e69de29..cb6356a 100644
--- a/pymic/__init__.py
+++ b/pymic/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+__version__ = "0.4.0"
\ No newline at end of file
diff --git a/pymic/io/__init__.py b/pymic/io/__init__.py
index e69de29..7f94f61 100644
--- a/pymic/io/__init__.py
+++ b/pymic/io/__init__.py
@@ -0,0 +1,4 @@
+from __future__ import absolute_import
+from pymic.io.image_read_write import *
+from pymic.io.nifty_dataset import *
+from pymic.io.h5_dataset import *
\ No newline at end of file
diff --git a/pymic/layer/__init__.py b/pymic/layer/__init__.py
index e69de29..72b8078 100644
--- a/pymic/layer/__init__.py
+++ b/pymic/layer/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/loss/__init__.py b/pymic/loss/__init__.py
index e69de29..72b8078 100644
--- a/pymic/loss/__init__.py
+++ b/pymic/loss/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/loss/cls/__init__.py b/pymic/loss/cls/__init__.py
index e69de29..72b8078 100644
--- a/pymic/loss/cls/__init__.py
+++ b/pymic/loss/cls/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/loss/cls/basic.py b/pymic/loss/cls/basic.py
index 4c90943..56925fc 100644
--- a/pymic/loss/cls/basic.py
+++ b/pymic/loss/cls/basic.py
@@ -65,10 +65,7 @@ def forward(self, loss_input_dict):
         labels  = loss_input_dict['ground_truth'][:, None] # reshape to N, 1
         softmax = nn.Softmax(dim = 1)
         predict = softmax(predict)
-        num_class  = list(predict.size())[1]
-        data_type = 'float' if(predict.dtype is torch.float32) else 'double'
-        soft_y = get_soft_label(labels, num_class, data_type)
-        loss = self.l1_loss(predict, soft_y)
+        loss = self.l1_loss(predict, labels)
         return loss
 
 class MSELoss(AbstractClassificationLoss):
@@ -84,10 +81,7 @@ def forward(self, loss_input_dict):
         labels  = loss_input_dict['ground_truth'][:, None] # reshape to N, 1
         softmax = nn.Softmax(dim = 1)
         predict = softmax(predict)
-        num_class  = list(predict.size())[1]
-        data_type = 'float' if(predict.dtype is torch.float32) else 'double'
-        soft_y = get_soft_label(labels, num_class, data_type)
-        loss = self.mse_loss(predict, soft_y)
+        loss = self.mse_loss(predict, labels)
         return loss
 
 class NLLLoss(AbstractClassificationLoss):
diff --git a/pymic/loss/seg/__init__.py b/pymic/loss/seg/__init__.py
index e69de29..72b8078 100644
--- a/pymic/loss/seg/__init__.py
+++ b/pymic/loss/seg/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/loss/seg/ce.py b/pymic/loss/seg/ce.py
index bbe6d02..529482b 100644
--- a/pymic/loss/seg/ce.py
+++ b/pymic/loss/seg/ce.py
@@ -18,7 +18,6 @@ class CrossEntropyLoss(AbstractSegLoss):
     """
     def __init__(self, params = None):
         super(CrossEntropyLoss, self).__init__(params)
-        
     
     def forward(self, loss_input_dict):
         predict = loss_input_dict['prediction']
diff --git a/pymic/loss/seg/deep_sup.py b/pymic/loss/seg/deep_sup.py
index 42ce172..da6d9ef 100644
--- a/pymic/loss/seg/deep_sup.py
+++ b/pymic/loss/seg/deep_sup.py
@@ -2,8 +2,43 @@
 from __future__ import print_function, division
 
 import torch.nn as nn
+from torch.nn.functional import interpolate
 from pymic.loss.seg.abstract import AbstractSegLoss
 
+def match_prediction_and_gt_shape(pred, gt, mode = 0):
+    pred_shape = list(pred.shape)
+    gt_shape   = list(gt.shape)
+    dim = len(pred_shape) - 2
+    shape_match = False 
+    if(dim == 2):
+        if(pred_shape[-1] == gt_shape[-1] and pred_shape[-2] == gt_shape[-2]):
+            shape_match = True
+    else:
+        if(pred_shape[-1] == gt_shape[-1] and pred_shape[-2] == gt_shape[-2]
+           and pred_shape[-3] == gt_shape[-3]):
+            shape_match = True
+    if(shape_match):
+        return pred, gt 
+    
+    interp_mode = 'bilinear' if dim == 2 else 'trilinear'
+    if(mode == 0):
+        pred_new = interpolate(pred, gt_shape[2:], mode = interp_mode)
+        gt_new   = gt  
+    elif(mode == 1):
+        pred_new = pred
+        gt_new   = interpolate(gt, pred_shape[2:], mode = interp_mode)
+    elif(mode == 2):
+        pred_new = pred
+        if(dim == 2):
+            avg_pool = nn.AdaptiveAvgPool2d(pred_shape[-2:])
+        else:
+            avg_pool = nn.AdaptiveAvgPool3d(pred_shape[-3:])
+        gt_new = avg_pool(gt)
+    else:
+        raise ValueError("mode shoud be 0, 1 or 2, but {0:} was given".format(mode))
+    return pred_new, gt_new
+            
+
 class DeepSuperviseLoss(AbstractSegLoss):
     '''
     Combine deep supervision with a basic loss function.  
@@ -12,28 +47,36 @@ class DeepSuperviseLoss(AbstractSegLoss):
 
     :param `loss_softmax`: (optional, bool) 
         Apply softmax to the prediction of network or not. Default is True.
-    :param `deep_suervise_weight`: (list) A list of weight for each deep supervision scale. \n
     :param `base_loss`: (nn.Module) The basic function used for each scale.
+    :param `deep_supervise_weight`: (list) A list of weight for each deep supervision scale. 
+    :param `deep_supervise_model`: (int) Mode for deep supervision when the prediction
+        has a smaller shape than the ground truth. 0: upsample the prediction to the size 
+        of the ground truth. 1: downsample the ground truth to the size of the prediction
+        via interpolation. 2: downsample the ground truth via adaptive average pooling.
 
     '''
     def __init__(self, params):
         super(DeepSuperviseLoss, self).__init__(params)
-        self.deep_sup_weight = params.get('deep_suervise_weight', None)
-        self.base_loss = params['base_loss']
+        self.base_loss       = params['base_loss']
+        self.deep_sup_weight = params.get('deep_supervise_weight', None)
+        self.deep_sup_mode   = params.get('deep_supervise_mode', 0)
 
     def forward(self, loss_input_dict):
-        predict = loss_input_dict['prediction']
-        if(not isinstance(predict, (list,tuple))):
+        pred = loss_input_dict['prediction']
+        gt   = loss_input_dict['ground_truth']
+        if(not isinstance(pred, (list,tuple))):
             raise ValueError("""For deep supervision, the prediction should
                 be a list or a tuple""")
-        predict_num = len(predict)
+        pred_num = len(pred)
         if(self.deep_sup_weight is None):
-            self.deep_sup_weight = [1.0] * predict_num
+            self.deep_sup_weight = [1.0] * pred_num
         else:
-            assert(predict_num == len(self.deep_sup_weight))
+            assert(pred_num == len(self.deep_sup_weight))
         loss_sum, weight_sum  = 0.0, 0.0
-        for i in range(predict_num):
-            loss_input_dict['prediction'] =  predict[i]
+        for i in range(pred_num):
+            pred_i, gt_i = match_prediction_and_gt_shape(pred[i], gt, self.deep_sup_mode)
+            loss_input_dict['prediction']   = pred_i
+            loss_input_dict['ground_truth'] = gt_i
             temp_loss   = self.base_loss(loss_input_dict)
             loss_sum   += temp_loss * self.deep_sup_weight[i]
             weight_sum += self.deep_sup_weight[i]
diff --git a/pymic/loss/seg/ssl.py b/pymic/loss/seg/ssl.py
index f15fc60..0bf276f 100644
--- a/pymic/loss/seg/ssl.py
+++ b/pymic/loss/seg/ssl.py
@@ -6,8 +6,9 @@
 import torch.nn as nn
 import numpy as np 
 from pymic.loss.seg.util import reshape_tensor_to_2D
+from pymic.loss.seg.abstract import AbstractSegLoss
 
-class EntropyLoss(nn.Module):
+class EntropyLoss(AbstractSegLoss):
     """
     Entropy Minimization for segmentation tasks.
     The parameters should be written in the `params` dictionary, and it has the
@@ -43,7 +44,7 @@ def forward(self, loss_input_dict):
         avg_ent = torch.mean(entropy)
         return avg_ent
     
-class TotalVariationLoss(nn.Module):
+class TotalVariationLoss(AbstractSegLoss):
     """
     Total Variation Loss for segmentation tasks.
     The parameters should be written in the `params` dictionary, and it has the
diff --git a/pymic/net/cls/__init__.py b/pymic/net/cls/__init__.py
index e69de29..72b8078 100644
--- a/pymic/net/cls/__init__.py
+++ b/pymic/net/cls/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/net/cls/torch_pretrained_net.py b/pymic/net/cls/torch_pretrained_net.py
index edc5d9c..5017f72 100644
--- a/pymic/net/cls/torch_pretrained_net.py
+++ b/pymic/net/cls/torch_pretrained_net.py
@@ -75,7 +75,7 @@ def __init__(self, params):
     def get_parameters_to_update(self):
         if(self.update_mode == "all"):
             return self.net.parameters()
-        elif(self.update_layers == "last"):
+        elif(self.update_mode == "last"):
             params = self.net.fc.parameters()
             if(self.in_chns !=3):
                 # combining the two iterables into a single one 
@@ -119,7 +119,7 @@ def get_parameters_to_update(self):
             params = self.net.classifier[-1].parameters()
             if(self.in_chns !=3):
                 params = itertools.chain()
-                for pram in [self.net.classifier[-1].parameters(), self.net.net.features[0].parameters()]:
+                for pram in [self.net.classifier[-1].parameters(), self.net.features[0].parameters()]:
                     params = itertools.chain(params, pram)
             return  params
         else:
@@ -138,7 +138,7 @@ class MobileNetV2(BuiltInNet):
         as well as the first layer when `input_chns` is not 3.
     """
     def __init__(self, params):
-        super(MobileNetV2, self).__init__()
+        super(MobileNetV2, self).__init__(params)
         self.net = models.mobilenet_v2(pretrained = self.pretrain)
         
         # replace the last layer 
@@ -157,7 +157,7 @@ def get_parameters_to_update(self):
             params = self.net.classifier[-1].parameters()
             if(self.in_chns !=3):
                 params = itertools.chain()
-                for pram in [self.net.classifier[-1].parameters(), self.net.net.features[0][0].parameters()]:
+                for pram in [self.net.classifier[-1].parameters(), self.net.features[0][0].parameters()]:
                     params = itertools.chain(params, pram)
             return  params
         else:
diff --git a/pymic/net/net2d/__init__.py b/pymic/net/net2d/__init__.py
index e69de29..72b8078 100644
--- a/pymic/net/net2d/__init__.py
+++ b/pymic/net/net2d/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/net/net2d/unet2d.py b/pymic/net/net2d/unet2d.py
index 2708800..9acc0ad 100644
--- a/pymic/net/net2d/unet2d.py
+++ b/pymic/net/net2d/unet2d.py
@@ -187,7 +187,7 @@ class UNet2D(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 deep_supervise: (bool) Using deep supervision for training or not.
+    :param multiscale_pred: (bool) Get multiscale prediction.
     """
     def __init__(self, params):
         super(UNet2D, self).__init__()
@@ -197,7 +197,7 @@ def __init__(self, params):
         self.dropout   = self.params['dropout']
         self.n_class   = self.params['class_num']
         self.bilinear  = self.params['bilinear']
-        self.deep_sup  = self.params['deep_supervise']
+        self.mul_pred  = self.params['multiscale_pred']
 
         assert(len(self.ft_chns) == 5 or len(self.ft_chns) == 4)
 
@@ -213,7 +213,7 @@ def __init__(self, params):
         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.deep_sup):
+        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)
@@ -239,19 +239,15 @@ 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.deep_sup):
-            out_shape = list(output.shape)[2:]
+        if(self.mul_pred):
             output1 = self.out_conv1(x_d1)
-            output1 = interpolate(output1, out_shape, mode = 'bilinear')
             output2 = self.out_conv2(x_d2)
-            output2 = interpolate(output2, out_shape, mode = 'bilinear')
             output3 = self.out_conv3(x_d3)
-            output3 = interpolate(output3, out_shape, mode = 'bilinear')
             output = [output, output1, output2, output3]
 
             if(len(x_shape) == 5):
-                new_shape = [N, D] + list(output[0].shape)[1:]
                 for i in range(len(output)):
+                    new_shape = [N, D] + list(output[i].shape)[1:]
                     output[i] = torch.transpose(torch.reshape(output[i], new_shape), 1, 2)
         elif(len(x_shape) == 5):
             new_shape = [N, D] + list(output.shape)[1:]
@@ -265,7 +261,8 @@ def forward(self, x):
               'feature_chns':[2, 8, 32, 48, 64],
               'dropout':  [0, 0, 0.3, 0.4, 0.5],
               'class_num': 2,
-              'bilinear': True}
+              'bilinear': True,
+              'multiscale_pred': False}
     Net = UNet2D(params)
     Net = Net.double()
 
diff --git a/pymic/net/net2d/unet2d_urpc.py b/pymic/net/net2d/unet2d_urpc.py
deleted file mode 100644
index ee8ab7c..0000000
--- a/pymic/net/net2d/unet2d_urpc.py
+++ /dev/null
@@ -1,132 +0,0 @@
-# -*- coding: utf-8 -*-
-from __future__ import print_function, division
-
-import torch
-import torch.nn as nn
-import numpy as np 
-from torch.distributions.uniform import Uniform
-from pymic.net.net2d.unet2d import ConvBlock, DownBlock, UpBlock
-
-def FeatureDropout(x):
-    attention = torch.mean(x, dim=1, keepdim=True)
-    max_val, _ = torch.max(attention.view(
-        x.size(0), -1), dim=1, keepdim=True)
-    threshold = max_val * np.random.uniform(0.7, 0.9)
-    threshold = threshold.view(x.size(0), 1, 1, 1).expand_as(attention)
-    drop_mask = (attention < threshold).float()
-    x = x.mul(drop_mask)
-    return x
-
-class FeatureNoise(nn.Module):
-    def __init__(self, uniform_range=0.3):
-        super(FeatureNoise, self).__init__()
-        self.uni_dist = Uniform(-uniform_range, uniform_range)
-
-    def feature_based_noise(self, x):
-        noise_vector = self.uni_dist.sample(
-            x.shape[1:]).to(x.device).unsqueeze(0)
-        x_noise = x.mul(noise_vector) + x
-        return x_noise
-
-    def forward(self, x):
-        x = self.feature_based_noise(x)
-        return x
-
-class UNet2D_URPC(nn.Module):
-    """
-    An modification the U-Net to obtain multi-scale prediction according to 
-    the URPC paper.
-
-    * Reference: Xiangde Luo, Guotai Wang*, Wenjun Liao, Jieneng Chen, Tao Song, Yinan Chen, 
-      Shichuan Zhang, Dimitris N. Metaxas, Shaoting Zhang. 
-      Semi-Supervised Medical Image Segmentation via Uncertainty Rectified Pyramid Consistency .
-      `Medical Image Analysis 2022. <https://doi.org/10.1016/j.media.2022.102517>`_
-
-    Also see: https://github.com/HiLab-git/SSL4MIS/blob/master/code/networks/unet.py
-    
-    Parameters are given in the `params` dictionary, and should include the
-    following fields:
-
-    :param in_chns: (int) Input channel number.
-    :param feature_chns: (list) Feature channel for each resolution level. 
-      The length should be 5, such as [16, 32, 64, 128, 256].
-    :param dropout: (list) The dropout ratio for each resolution level. 
-      The length should be the same as that of `feature_chns`.
-    :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.
-    """
-    def __init__(self, params):
-        super(UNet2D_URPC, self).__init__()
-        self.params    = params
-        self.in_chns   = self.params['in_chns']
-        self.ft_chns   = self.params['feature_chns']
-        self.dropout   = self.params['dropout']
-        self.n_class   = self.params['class_num']
-        self.bilinear  = self.params['bilinear']
-        assert(len(self.ft_chns) == 5)
-
-        self.in_conv= ConvBlock(self.in_chns, self.ft_chns[0], self.dropout[0])
-        self.down1  = DownBlock(self.ft_chns[0], self.ft_chns[1], self.dropout[1])
-        self.down2  = DownBlock(self.ft_chns[1], self.ft_chns[2], self.dropout[2])
-        self.down3  = DownBlock(self.ft_chns[2], self.ft_chns[3], self.dropout[3])
-        self.down4  = DownBlock(self.ft_chns[3], self.ft_chns[4], self.dropout[4])
-        self.up1 = UpBlock(self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], 0.0, self.bilinear) 
-        self.up2 = UpBlock(self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], 0.0, self.bilinear) 
-        self.up3 = UpBlock(self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], 0.0, self.bilinear) 
-        self.up4 = UpBlock(self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], 0.0, self.bilinear) 
-    
-        self.out_conv = nn.Conv2d(self.ft_chns[0], self.n_class,  
-                                  kernel_size = 3, padding = 1)
-        self.out_conv_dp4 = nn.Conv2d(self.ft_chns[4], self.n_class,
-                                      kernel_size=3, padding=1)
-        self.out_conv_dp3 = nn.Conv2d(self.ft_chns[3], self.n_class,
-                                      kernel_size=3, padding=1)
-        self.out_conv_dp2 = nn.Conv2d(self.ft_chns[2], self.n_class,
-                                      kernel_size=3, padding=1)
-        self.out_conv_dp1 = nn.Conv2d(self.ft_chns[1], self.n_class,
-                                      kernel_size=3, padding=1)
-        self.feature_noise = FeatureNoise()
-
-    def forward(self, x):
-        x_shape = list(x.shape)
-        if(len(x_shape) == 5):
-          [N, C, D, H, W] = x_shape
-          new_shape = [N*D, C, H, W]
-          x = torch.transpose(x, 1, 2)
-          x = torch.reshape(x, new_shape)
-        x0 = self.in_conv(x)
-        x1 = self.down1(x0)
-        x2 = self.down2(x1)
-        x3 = self.down3(x2)
-        x4 = self.down4(x3)
-        
-        x = self.up1(x4, x3)
-        if self.training:
-            x = nn.functional.dropout(x, p=0.5)
-        dp3_out = self.out_conv_dp3(x)
-
-        x = self.up2(x, x2)
-        if self.training:
-            x = FeatureDropout(x)
-        dp2_out = self.out_conv_dp2(x)
-
-        x = self.up3(x, x1)
-        if self.training:
-            x = self.feature_noise(x)
-        dp1_out = self.out_conv_dp1(x)
-
-        x = self.up4(x, x0)
-        dp0_out = self.out_conv(x)
-
-        out_shape = list(dp0_out.shape)[2:]
-        dp3_out = nn.functional.interpolate(dp3_out, out_shape)
-        dp2_out = nn.functional.interpolate(dp2_out, out_shape)
-        dp1_out = nn.functional.interpolate(dp1_out, out_shape)
-        out = [dp0_out, dp1_out, dp2_out, dp3_out]
-
-        if(len(x_shape) == 5):
-            new_shape = [N, D] + list(dp0_out.shape)[1:]
-            for i in range(len(out)):
-                out[i] = torch.transpose(torch.reshape(out[i], new_shape), 1, 2)
-        return out
\ No newline at end of file
diff --git a/pymic/net/net3d/__init__.py b/pymic/net/net3d/__init__.py
index e69de29..72b8078 100644
--- a/pymic/net/net3d/__init__.py
+++ b/pymic/net/net3d/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/net/net3d/unet3d.py b/pymic/net/net3d/unet3d.py
index af28d73..a17bcb8 100644
--- a/pymic/net/net3d/unet3d.py
+++ b/pymic/net/net3d/unet3d.py
@@ -77,6 +77,104 @@ def forward(self, x1, x2):
         x = torch.cat([x2, x1], dim=1)
         return self.conv(x)
 
+class Encoder(nn.Module):
+    """
+    Encoder of 3D UNet.
+
+    Parameters are given in the `params` dictionary, and should include the
+    following fields:
+
+    :param in_chns: (int) Input channel number.
+    :param feature_chns: (list) Feature channel for each resolution level. 
+      The length should be 4 or 5, such as [16, 32, 64, 128, 256].
+    :param dropout: (list) The dropout ratio for each resolution level. 
+      The length should be the same as that of `feature_chns`.
+    """
+    def __init__(self, params):
+        super(Encoder, self).__init__()
+        self.params    = params
+        self.in_chns   = self.params['in_chns']
+        self.ft_chns   = self.params['feature_chns']
+        self.dropout   = self.params['dropout']
+        assert(len(self.ft_chns) == 5 or len(self.ft_chns) == 4)
+
+        self.in_conv= ConvBlock(self.in_chns, self.ft_chns[0], self.dropout[0])
+        self.down1  = DownBlock(self.ft_chns[0], self.ft_chns[1], self.dropout[1])
+        self.down2  = DownBlock(self.ft_chns[1], self.ft_chns[2], self.dropout[2])
+        self.down3  = DownBlock(self.ft_chns[2], self.ft_chns[3], self.dropout[3])
+        if(len(self.ft_chns) == 5):
+            self.down4  = DownBlock(self.ft_chns[3], self.ft_chns[4], self.dropout[4])
+
+    def forward(self, x):
+        x0 = self.in_conv(x)
+        x1 = self.down1(x0)
+        x2 = self.down2(x1)
+        x3 = self.down3(x2)
+        output = [x0, x1, x2, x3]
+        if(len(self.ft_chns) == 5):
+          x4 = self.down4(x3)
+          output.append(x4)
+        return output
+
+class Decoder(nn.Module):
+    """
+    Decoder of 3D UNet.
+
+    Parameters are given in the `params` dictionary, and should include the
+    following fields:
+
+    :param in_chns: (int) Input channel number.
+    :param feature_chns: (list) Feature channel for each resolution level. 
+      The length should be 4 or 5, such as [16, 32, 64, 128, 256].
+    :param dropout: (list) The dropout ratio for each resolution level. 
+      The length should be the same as that of `feature_chns`.
+    :param class_num: (int) The class number for segmentation task. 
+    :param trilinear: (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__()
+        self.params    = params
+        self.in_chns   = self.params['in_chns']
+        self.ft_chns   = self.params['feature_chns']
+        self.dropout   = self.params['dropout']
+        self.n_class   = self.params['class_num']
+        self.trilinear = self.params.get('trilinear', True)
+        self.mul_pred  = self.params.get('multiscale_pred', False)
+
+        assert(len(self.ft_chns) == 5 or len(self.ft_chns) == 4)
+
+        if(len(self.ft_chns) == 5):
+            self.up1 = UpBlock(self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], self.dropout[3], self.trilinear) 
+        self.up2 = UpBlock(self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], self.dropout[2], self.trilinear) 
+        self.up3 = UpBlock(self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], self.dropout[1], self.trilinear) 
+        self.up4 = UpBlock(self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], self.dropout[0], self.trilinear) 
+        self.out_conv = nn.Conv3d(self.ft_chns[0], self.n_class, kernel_size = 1)
+        if(self.mul_pred):
+            self.out_conv1 = nn.Conv3d(self.ft_chns[1], self.n_class, kernel_size = 1)
+            self.out_conv2 = nn.Conv3d(self.ft_chns[2], self.n_class, kernel_size = 1)
+            self.out_conv3 = nn.Conv3d(self.ft_chns[3], self.n_class, kernel_size = 1)
+
+    def forward(self, x):
+        if(len(self.ft_chns) == 5):
+            assert(len(x) == 5)
+            x0, x1, x2, x3, x4 = x 
+            x_d3 = self.up1(x4, x3)
+        else:
+            assert(len(x) == 4)
+            x0, x1, x2, x3 = x 
+            x_d3 = x3
+        x_d2 = self.up2(x_d3, x2)
+        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 UNet3D(nn.Module):
     """
@@ -100,7 +198,7 @@ class UNet3D(nn.Module):
     :param class_num: (int) The class number for segmentation task. 
     :param trilinear: (bool) Using trilinear for up-sampling or not. 
         If False, deconvolution will be used for up-sampling.
-    :param deep_supervise: (bool) Using deep supervision for training or not.
+    :param multiscale_pred: (bool) Get multi-scale prediction.
     """
     def __init__(self, params):
         super(UNet3D, self).__init__()
@@ -110,7 +208,7 @@ def __init__(self, params):
         self.dropout   = self.params['dropout']
         self.n_class   = self.params['class_num']
         self.trilinear = self.params['trilinear']
-        self.deep_sup  = self.params['deep_supervise']
+        self.mul_pred  = self.params['multiscale_pred']
         assert(len(self.ft_chns) == 5 or len(self.ft_chns) == 4)
 
         self.in_conv= ConvBlock(self.in_chns, self.ft_chns[0], self.dropout[0])
@@ -129,7 +227,7 @@ def __init__(self, params):
                dropout_p = self.dropout[0], trilinear=self.trilinear) 
     
         self.out_conv = nn.Conv3d(self.ft_chns[0], self.n_class, kernel_size = 1)
-        if(self.deep_sup):
+        if(self.mul_pred):
             self.out_conv1 = nn.Conv3d(self.ft_chns[1], self.n_class, kernel_size = 1)
             self.out_conv2 = nn.Conv3d(self.ft_chns[2], self.n_class, kernel_size = 1)
             self.out_conv3 = nn.Conv3d(self.ft_chns[3], self.n_class, kernel_size = 1)
@@ -148,14 +246,10 @@ 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.deep_sup):
-            out_shape = list(output.shape)[2:]
+        if(self.mul_pred):
             output1 = self.out_conv1(x_d1)
-            output1 = interpolate(output1, out_shape, mode = 'trilinear')
             output2 = self.out_conv2(x_d2)
-            output2 = interpolate(output2, out_shape, mode = 'trilinear')
             output3 = self.out_conv3(x_d3)
-            output3 = interpolate(output3, out_shape, mode = 'trilinear')
             output = [output, output1, output2, output3]
         return output
 
@@ -164,7 +258,8 @@ def forward(self, x):
               'class_num': 2,
               'feature_chns':[2, 8, 32, 64],
               'dropout' : [0, 0, 0, 0.5],
-              'trilinear': True}
+              'trilinear': True,
+              'multiscale_pred': False}
     Net = UNet3D(params)
     Net = Net.double()
 
diff --git a/pymic/net/net3d/unet3d_dual_branch.py b/pymic/net/net3d/unet3d_dual_branch.py
new file mode 100644
index 0000000..3bede4e
--- /dev/null
+++ b/pymic/net/net3d/unet3d_dual_branch.py
@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+from pymic.net.net3d.unet3d import *
+
+class UNet3D_DualBranch(nn.Module):
+    """
+    A dual branch network using UNet3D as backbone.
+
+    * Reference: Xiangde Luo, Minhao Hu, Wenjun Liao, Shuwei Zhai, Tao Song, Guotai Wang,
+      Shaoting Zhang. ScribblScribble-Supervised Medical Image Segmentation via 
+      Dual-Branch Network and Dynamically Mixed Pseudo Labels Supervision.
+      `MICCAI 2022. <https://arxiv.org/abs/2203.02106>`_ 
+
+    The parameters for the backbone should be given in the `params` dictionary. 
+    See :mod:`pymic.net.net3d.unet3d.UNet3D` for details. 
+    In addition, the following field should be included:
+
+    :param output_mode: (str) How to obtain the result during the inference. 
+      `average`: taking average of the two branches. 
+      `first`: takeing the result in the first branch. 
+      `second`: taking the result in the second branch.
+    """
+    def __init__(self, params):
+        super(UNet3D_DualBranch, self).__init__()
+        self.output_mode = params.get("output_mode", "average")
+        self.encoder  = Encoder(params)
+        self.decoder1 = Decoder(params)    
+        self.decoder2 = Decoder(params)        
+
+    def forward(self, x):
+        f = self.encoder(x)
+        output1 = self.decoder1(f)
+        output2 = self.decoder2(f)
+
+        if(self.training):
+          return output1, output2
+        else:
+          if(self.output_mode == "average"):
+            return (output1 + output2)/2
+          elif(self.output_mode == "first"):
+            return output1
+          else:
+            return output2
diff --git a/pymic/net/net_dict_seg.py b/pymic/net/net_dict_seg.py
index 5dd3bee..fc7692f 100644
--- a/pymic/net/net_dict_seg.py
+++ b/pymic/net/net_dict_seg.py
@@ -4,7 +4,6 @@
 
 * UNet2D :mod:`pymic.net.net2d.unet2d.UNet2D`
 * UNet2D_DualBranch :mod:`pymic.net.net2d.unet2d_dual_branch.UNet2D_DualBranch`
-* UNet2D_URPC :mod:`pymic.net.net2d.unet2d_urpc.UNet2D_URPC`
 * UNet2D_CCT  :mod:`pymic.net.net2d.unet2d_cct.UNet2D_CCT`
 * UNet2D_ScSE  :mod:`pymic.net.net2d.unet2d_scse.UNet2D_ScSE`
 * AttentionUNet2D  :mod:`pymic.net.net2d.unet2d_attention.AttentionUNet2D`
@@ -17,7 +16,6 @@
 from __future__ import print_function, division
 from pymic.net.net2d.unet2d import UNet2D
 from pymic.net.net2d.unet2d_dual_branch import UNet2D_DualBranch
-from pymic.net.net2d.unet2d_urpc import UNet2D_URPC
 from pymic.net.net2d.unet2d_cct import UNet2D_CCT
 from pymic.net.net2d.cople_net import COPLENet
 from pymic.net.net2d.unet2d_attention import AttentionUNet2D
@@ -26,11 +24,11 @@
 from pymic.net.net3d.unet2d5 import UNet2D5
 from pymic.net.net3d.unet3d import UNet3D
 from pymic.net.net3d.unet3d_scse import UNet3D_ScSE
+from pymic.net.net3d.unet3d_dual_branch import UNet3D_DualBranch
 
 SegNetDict = {
 	'UNet2D': UNet2D,
 	'UNet2D_DualBranch': UNet2D_DualBranch,
-	'UNet2D_URPC': UNet2D_URPC,
 	'UNet2D_CCT': UNet2D_CCT,
 	'COPLENet': COPLENet,
 	'AttentionUNet2D': AttentionUNet2D,
@@ -38,5 +36,7 @@
 	'UNet2D_ScSE': UNet2D_ScSE,
 	'UNet2D5': UNet2D5,
 	'UNet3D': UNet3D,
-	'UNet3D_ScSE': UNet3D_ScSE
+	'UNet3D_ScSE': UNet3D_ScSE,
+	'UNet3D_DualBranch': UNet3D_DualBranch
+
 	}
diff --git a/pymic/net_run/__init__.py b/pymic/net_run/__init__.py
index e69de29..72b8078 100644
--- a/pymic/net_run/__init__.py
+++ b/pymic/net_run/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/net_run/agent_abstract.py b/pymic/net_run/agent_abstract.py
index dc6b1ad..9131ba0 100644
--- a/pymic/net_run/agent_abstract.py
+++ b/pymic/net_run/agent_abstract.py
@@ -259,7 +259,7 @@ def worker_init_fn(worker_id):
 
             bn_train = self.config['dataset']['train_batch_size']
             bn_valid = self.config['dataset'].get('valid_batch_size', 1)
-            num_worker = self.config['dataset'].get('num_workder', 16)
+            num_worker = self.config['dataset'].get('num_worker', 16)
             g_train, g_valid = torch.Generator(), torch.Generator()
             g_train.manual_seed(self.random_seed)
             g_valid.manual_seed(self.random_seed)
@@ -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)
diff --git a/pymic/net_run/agent_cls.py b/pymic/net_run/agent_cls.py
index bb04a9d..5610982 100644
--- a/pymic/net_run/agent_cls.py
+++ b/pymic/net_run/agent_cls.py
@@ -5,17 +5,20 @@
 import csv
 import logging
 import time
-import torch
-from torchvision import transforms
 import numpy as np
+import torch
 import torch.nn as nn
 from datetime import datetime
+from random import random
+from torch.optim import lr_scheduler
+from torchvision import transforms
 from tensorboardX import SummaryWriter
 from pymic.io.nifty_dataset import ClassificationDataset
 from pymic.loss.loss_dict_cls import PyMICClsLossDict
 from pymic.net.net_dict_cls import TorchClsNetDict
 from pymic.transform.trans_dict import TransformDict
 from pymic.net_run.agent_abstract import NetRunAgent
+from pymic.util.general import mixup
 import warnings
 warnings.filterwarnings('ignore', '.*output shape of zoom.*')
 
@@ -110,16 +113,17 @@ def get_evaluation_score(self, outputs, labels):
         """
         Get evaluation score for a prediction.
 
-        :param outputs: (tensor) Prediction obtained by a network. 
-        :param labels: (tensor) The ground truth.
+        :param outputs: (tensor) Prediction obtained by a network with size N X C. 
+        :param labels: (tensor) The ground truth with size N X C.
         """
         metrics = self.config['training'].get("evaluation_metric", "accuracy")
         if(metrics != "accuracy"): # default classification accuracy
             raise ValueError("Not implemeted for metric {0:}".format(metrics))
         if(self.task_type == "cls"):
-            _, preds = torch.max(outputs, 1)
-            consis= self.convert_tensor_type(preds ==  labels.data)
-            score = torch.mean(consis)
+            out_argmax = torch.argmax(outputs, 1)
+            lab_argmax = torch.argmax(labels, 1)
+            consis = self.convert_tensor_type(out_argmax ==  lab_argmax)
+            score  = torch.mean(consis) 
         elif(self.task_type == "cls_nexcl"): #nonexclusive classification
             preds = self.convert_tensor_type(outputs > 0.5)
             consis= self.convert_tensor_type(preds ==  labels.data)
@@ -128,6 +132,7 @@ def get_evaluation_score(self, outputs, labels):
 
     def training(self):
         iter_valid   = self.config['training']['iter_valid']
+        mixup_prob   = self.config['training'].get('mixup_probability', 0.5)
         sample_num   = 0
         running_loss = 0
         running_score= 0
@@ -139,16 +144,19 @@ def training(self):
                 self.trainIter = iter(self.train_loader)
                 data = next(self.trainIter)
             inputs = self.convert_tensor_type(data['image'])
-            labels = data['label'].long()         
+            labels = self.convert_tensor_type(data['label_prob'])  
+            if(random() < mixup_prob):
+                inputs, labels = mixup(inputs, labels)    
             inputs, labels = inputs.to(self.device), labels.to(self.device)
+
             # zero the parameter gradients
             self.optimizer.zero_grad()
             # forward + backward + optimize
             outputs = self.net(inputs)
-            loss = self.get_loss_value(data, inputs, outputs, labels)
+            
+            loss = self.get_loss_value(data, outputs, labels)
             loss.backward()
             self.optimizer.step()
-            self.scheduler.step()
             
             # statistics
             sample_num   += labels.size(0)
@@ -170,12 +178,12 @@ def validation(self):
             self.net.eval()
             for data in validIter:
                 inputs = self.convert_tensor_type(data['image'])
-                labels = data['label'].long()             
+                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, inputs, outputs, labels)
+                loss = self.get_loss_value(data, outputs, labels)
                                 
                 # statistics
                 sample_num   += labels.size(0)
@@ -184,19 +192,18 @@ def validation(self):
 
         avg_loss = running_loss / sample_num
         avg_score= running_score.double() / sample_num
-        metrics =self.config['training'].get("evaluation_metric", "accuracy")
+        metrics  = self.config['training'].get("evaluation_metric", "accuracy")
         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(
@@ -221,7 +228,15 @@ def train_valid(self):
         iter_max    = self.config['training']['iter_max']
         iter_valid  = self.config['training']['iter_valid']
         iter_save   = self.config['training']['iter_save']
+        early_stop_it = self.config['training'].get('early_stop_patience', None)
         metrics     = self.config['training'].get("evaluation_metric", "accuracy")
+        if(iter_save is None):
+            iter_save_list = [iter_max]
+        elif(isinstance(iter_save, (tuple, list))):
+            iter_save_list = iter_save
+        else:
+            iter_save_list = range(0, iter_max + 1, iter_save)
+        
         self.max_val_score  = 0.0
         self.max_val_it     = 0
         self.best_model_wts = None 
@@ -230,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']
@@ -242,28 +260,48 @@ def train_valid(self):
 
         logging.info("{0:} training start".format(str(datetime.now())[:-7]))
         self.summ_writer = SummaryWriter(self.config['training']['ckpt_save_dir'])
+        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()
-            glob_it = it + iter_valid
-            self.write_scalars(train_scalars, valid_scalars, glob_it)
+            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    = glob_it
-                self.best_model_wts = copy.deepcopy(self.net.state_dict())
+                self.max_val_it    = self.glob_it
+                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())
             
-            if (glob_it % iter_save ==  0):
-                save_dict = {'iteration': glob_it,
+            stop_now = True if(early_stop_it is not None and \
+                self.glob_it - self.max_val_it > early_stop_it) else False
+
+            if ((self.glob_it in iter_save_list) or stop_now):
+                save_dict = {'iteration': self.glob_it,
                              'valid_pred': valid_scalars[metrics],
                              'model_state_dict': self.net.state_dict(),
                              'optimizer_state_dict': self.optimizer.state_dict()}
-                save_name = "{0:}/{1:}_{2:}.pt".format(ckpt_dir, ckpt_prefix, glob_it)
+                save_name = "{0:}/{1:}_{2:}.pt".format(ckpt_dir, ckpt_prefix, self.glob_it)
                 torch.save(save_dict, save_name) 
                 txt_file = open("{0:}/{1:}_latest.txt".format(ckpt_dir, ckpt_prefix), 'wt')
-                txt_file.write(str(glob_it))
+                txt_file.write(str(self.glob_it))
                 txt_file.close()
-
+            if(stop_now):
+                logging.info("The training is early stopped")
+                break
         # save the best performing checkpoint
         save_dict = {'iteration': self.max_val_it,
                     'valid_pred': self.max_val_score,
@@ -278,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]))
@@ -290,8 +327,8 @@ def infer(self):
         
         if(self.config['testing'].get('evaluation_mode', True)):
             self.net.eval()
-            
-        output_csv   = self.config['testing']['output_csv']
+        
+        output_csv   = self.config['testing']['output_dir'] + '/' + self.config['testing']['output_csv']
         class_num    = self.config['network']['class_num']
         save_probability = self.config['testing'].get('save_probability', False)
         
@@ -327,7 +364,7 @@ def infer(self):
             csv_writer = csv.writer(csv_file, delimiter=',', 
                                 quotechar='"',quoting=csv.QUOTE_MINIMAL)
             head = ['image', 'label']
-            if(len(out_lab_list[0]) > 1):
+            if(len(out_lab_list[0]) > 2):
                 head = ['image'] + ['label{0:}'.format(i) for i in range(class_num)]
             csv_writer.writerow(head)
             for item in out_lab_list:
diff --git a/pymic/net_run/agent_seg.py b/pymic/net_run/agent_seg.py
index bba4508..887a516 100644
--- a/pymic/net_run/agent_seg.py
+++ b/pymic/net_run/agent_seg.py
@@ -10,9 +10,10 @@
 import numpy as np
 import torch.nn as nn
 import torch.optim as optim
-from torch.optim import lr_scheduler
 import torch.nn.functional as F
 from datetime import datetime
+from random import random
+from torch.optim import lr_scheduler
 from tensorboardX import SummaryWriter
 from pymic.io.image_read_write import save_nd_array_as_image
 from pymic.io.nifty_dataset import NiftyDataset
@@ -28,6 +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, tensor_shape_match
 
 class SegmentationAgent(NetRunAgent):
     def __init__(self, config, stage = 'train'):
@@ -94,9 +96,12 @@ 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
@@ -104,7 +109,10 @@ def create_loss_calculator(self):
     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
     
@@ -120,6 +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.0)
         train_loss  = 0
         train_dice_list = []
         self.net.train()
@@ -132,7 +141,9 @@ def training(self):
             # get the inputs
             inputs      = self.convert_tensor_type(data['image'])
             labels_prob = self.convert_tensor_type(data['label_prob'])                 
-            
+            if(mixup_prob > 0 and random() < mixup_prob):
+                inputs, labels_prob = mixup(inputs, labels_prob) 
+                   
             # # for debug
             # for i in range(inputs.shape[0]):
             #     image_i = inputs[i][0]
@@ -157,10 +168,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)):
@@ -172,9 +179,9 @@ def training(self):
             train_dice_list.append(dice_list.cpu().numpy())
         train_avg_loss = train_loss / 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, 'avg_dice':train_avg_dice,\
+        train_scalers = {'loss': train_avg_loss, 'avg_fg_dice':train_avg_dice,\
             'class_dice': train_cls_dice}
         return train_scalers
         
@@ -213,18 +220,14 @@ 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,\
+        valid_avg_dice = valid_cls_dice[1:].mean()
+        valid_scalers = {'loss': valid_avg_loss, 'avg_fg_dice': valid_avg_dice,\
             'class_dice': valid_cls_dice}
         return valid_scalers
 
     def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
         loss_scalar ={'train':train_scalars['loss'], 'valid':valid_scalars['loss']}
-        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('dice', dice_scalar, glob_it)
         self.summ_writer.add_scalars('lr', {"lr": lr_value}, glob_it)
@@ -234,11 +237,11 @@ 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']) + "]")        
 
     def train_valid(self):
@@ -249,11 +252,15 @@ 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]
-        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)
@@ -268,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)
@@ -295,16 +313,20 @@ 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_fg_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))
             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['avg_dice'] > self.max_val_dice):
-                self.max_val_dice = valid_scalars['avg_dice']
+            if(valid_scalars['avg_fg_dice'] > self.max_val_dice):
+                self.max_val_dice = valid_scalars['avg_fg_dice']
                 self.max_val_it   = self.glob_it
                 if(len(device_ids) > 1):
                     self.best_model_wts = copy.deepcopy(self.net.module.state_dict())
@@ -315,7 +337,7 @@ def train_valid(self):
                 self.glob_it - self.max_val_it > early_stop_it) else False
             if ((self.glob_it in iter_save_list) or stop_now):
                 save_dict = {'iteration': self.glob_it,
-                             'valid_pred': valid_scalars['avg_dice'],
+                             'valid_pred': valid_scalars['avg_fg_dice'],
                              'model_state_dict': self.net.module.state_dict() \
                                  if len(device_ids) > 1 else self.net.state_dict(),
                              'optimizer_state_dict': self.optimizer.state_dict()}
diff --git a/pymic/net_run/get_optimizer.py b/pymic/net_run/get_optimizer.py
index c4504de..ad8fda0 100644
--- a/pymic/net_run/get_optimizer.py
+++ b/pymic/net_run/get_optimizer.py
@@ -10,48 +10,69 @@ def get_optimizer(name, net_params, optim_params):
     lr = optim_params['learning_rate']
     momentum = optim_params['momentum']
     weight_decay = optim_params['weight_decay']
+    # see https://www.codeleading.com/article/44815584159/
+    param_group = [{'params': net_params, 'initial_lr': lr}]
     if(keyword_match(name, "SGD")):
-        return optim.SGD(net_params, lr, 
+        return optim.SGD(param_group, lr, 
             momentum = momentum, weight_decay = weight_decay)
     elif(keyword_match(name, "Adam")):
-        return optim.Adam(net_params, lr, weight_decay = weight_decay)
+        return optim.Adam(param_group, lr, weight_decay = weight_decay)
     elif(keyword_match(name, "SparseAdam")):
-        return optim.SparseAdam(net_params, lr)
+        return optim.SparseAdam(param_group, lr)
     elif(keyword_match(name, "Adadelta")):
-        return optim.Adadelta(net_params, lr, weight_decay = weight_decay)
+        return optim.Adadelta(param_group, lr, weight_decay = weight_decay)
     elif(keyword_match(name, "Adagrad")): 
-        return optim.Adagrad(net_params, lr, weight_decay = weight_decay)
+        return optim.Adagrad(param_group, lr, weight_decay = weight_decay)
     elif(keyword_match(name, "Adamax")): 
-        return optim.Adamax(net_params, lr, weight_decay = weight_decay)
+        return optim.Adamax(param_group, lr, weight_decay = weight_decay)
     elif(keyword_match(name, "ASGD")): 
-        return optim.ASGD(net_params, lr, weight_decay = weight_decay)
+        return optim.ASGD(param_group, lr, weight_decay = weight_decay)
     elif(keyword_match(name, "LBFGS")): 
-        return optim.LBFGS(net_params, lr)
+        return optim.LBFGS(param_group, lr)
     elif(keyword_match(name, "RMSprop")): 
-        return optim.RMSprop(net_params, lr, momentum = momentum,
+        return optim.RMSprop(param_group, lr, momentum = momentum,
             weight_decay = weight_decay)
     elif(keyword_match(name, "Rprop")): 
-        return optim.Rprop(net_params, lr)
+        return optim.Rprop(param_group, lr)
     else:
         raise ValueError("unsupported optimizer {0:}".format(name))
 
 
 def get_lr_scheduler(optimizer, sched_params):
-    name = sched_params["lr_scheduler"]
+    name       = sched_params["lr_scheduler"]
+    val_it     = sched_params["iter_valid"]
+    epoch_last = sched_params["last_iter"] 
+    if(epoch_last > 0):
+        epoch_last = int(epoch_last / val_it)
     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]
+        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, "StepLR")):
+        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)
+    elif(keyword_match(name, "CosineAnnealingLR")):
+        epoch_max  = sched_params["iter_max"] / val_it
+        lr_min     = sched_params.get("lr_min", 0)
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
+                    epoch_max, lr_min, epoch_last)
+    elif(keyword_match(name, "PolynomialLR")):
+        epoch_max  = sched_params["iter_max"] / val_it
+        power      = sched_params["lr_power"]
+        scheduler = lr_scheduler.PolynomialLR(optimizer,
+                    epoch_max, power, epoch_last)
     else:
         raise ValueError("unsupported lr scheduler {0:}".format(name))
     return scheduler
\ No newline at end of file
diff --git a/pymic/net_run/infer_func.py b/pymic/net_run/infer_func.py
index 81d0b53..43162d0 100644
--- a/pymic/net_run/infer_func.py
+++ b/pymic/net_run/infer_func.py
@@ -15,7 +15,8 @@ class Inferer(object):
     :param `sliding_window_stride`: (optional, list) The sliding window stride. 
     :param `tta_mode`: (optional, int) The test time augmentation mode. Default
         is 0 (no test time augmentation). The other option is 1 (augmentation 
-        with horinzontal and vertical flipping).
+        with horinzontal and vertical flipping) and 2 (ensemble of inference
+        in axial, sagittal and coronal views for 2D networks applied to 3D volumes)
     """
     def __init__(self, config):
         self.config = config
@@ -170,6 +171,22 @@ def run(self, model, image):
                 outputs3 = torch.flip(outputs3, [-1])
                 outputs4 = torch.flip(outputs4, [-2, -1])
                 outputs = (outputs1 + outputs2 + outputs3 + outputs4) / 4
+        elif(tta_mode == 2):
+            outputs1 = self.__infer(image)
+            outputs2 = self.__infer(torch.transpose(image, -1, -3))
+            outputs3 = self.__infer(torch.transpose(image, -2, -3))
+            if(isinstance(outputs1, (tuple, list))):
+                outputs = []
+                for i in range(len(outputs1)):
+                    temp_out1 = outputs1[i]
+                    temp_out2 = torch.transpose(outputs2[i], -1, -3)
+                    temp_out3 = torch.transpose(outputs3[i], -2, -3)
+                    temp_mean = (temp_out1 + temp_out2 + temp_out3) / 3
+                    outputs.append(temp_mean)
+            else:
+                outputs2 = torch.transpose(outputs2, -1, -3)
+                outputs3 = torch.transpose(outputs3, -2, -3)
+                outputs = (outputs1 + outputs2 + outputs3) / 3
         else:
             raise ValueError("Undefined tta_mode {0:}".format(tta_mode))
         return outputs
diff --git a/pymic/net_run/noisy_label/__init__.py b/pymic/net_run/noisy_label/__init__.py
new file mode 100644
index 0000000..448a5e4
--- /dev/null
+++ b/pymic/net_run/noisy_label/__init__.py
@@ -0,0 +1,8 @@
+from __future__ import absolute_import
+from pymic.net_run.noisy_label.nll_co_teaching import NLLCoTeaching
+from pymic.net_run.noisy_label.nll_trinet import NLLTriNet
+from pymic.net_run.noisy_label.nll_dast import NLLDAST
+
+NLLMethodDict = {'CoTeaching': NLLCoTeaching,
+    "TriNet": NLLTriNet,
+    "DAST": NLLDAST}
\ No newline at end of file
diff --git a/pymic/net_run_nll/nll_clslsr.py b/pymic/net_run/noisy_label/nll_clslsr.py
similarity index 94%
rename from pymic/net_run_nll/nll_clslsr.py
rename to pymic/net_run/noisy_label/nll_clslsr.py
index 722b588..0148621 100644
--- a/pymic/net_run_nll/nll_clslsr.py
+++ b/pymic/net_run/noisy_label/nll_clslsr.py
@@ -3,7 +3,6 @@
 import logging
 import os
 import scipy
-import sys
 import torch
 import numpy as np 
 import pandas as pd
@@ -27,7 +26,11 @@ def get_confident_map(gt, pred, CL_type = 'both'):
 
     :return: A tensor representing the noisiness of each pixel.
     """
-    import cleanlab
+    try:
+        import cleanlab
+        assert(cleanlab.__version__ == '1.0.1')
+    except:
+        raise ValueError("Error: cleanlab 1.0.1 required. Please install it by `pip install cleanlab==1.0.1`")
     prob = scipy.special.softmax(pred, axis = 1)
     if CL_type in ['both', 'Qij']:
         noise = cleanlab.pruning.get_noise_indices(gt, prob, prune_method='both', n_jobs=1)
@@ -146,15 +149,7 @@ def test_time_dropout(m):
             dst_path = os.path.join(save_dir, filename)
             conf_map.save(dst_path)
 
-def get_confidence_map():
-    """
-    The main function to get the confidence map during inference.
-    """
-    if(len(sys.argv) < 2):
-        print('Number of arguments should be 3. e.g.')
-        print('   python nll_clslsr.py config.cfg')
-        exit()
-    cfg_file = str(sys.argv[1])
+def get_confidence_map(cfg_file):
     config   = parse_config(cfg_file)
     config   = synchronize_config(config)
 
@@ -173,7 +168,7 @@ def get_confidence_map():
             one_transform = transform_dict[name](transform_param)
             transform_list.append(one_transform)
         data_transform = transforms.Compose(transform_list)
-    print('transform list', transform_list)
+
     csv_file  = config['dataset']['train_csv']
     modal_num = config['dataset'].get('modal_num', 1)
     dataset  = NiftyDataset(root_dir  = config['dataset']['root_dir'],
@@ -201,7 +196,4 @@ def get_confidence_map():
                    "label": df_train["label"]}
     train_cl_csv = csv_file.replace(".csv", "_clslsr.csv")
     df_cl = pd.DataFrame.from_dict(train_cl_dict)
-    df_cl.to_csv(train_cl_csv, index = False)
-
-if __name__ == "__main__":
-    get_confidence_map()
\ No newline at end of file
+    df_cl.to_csv(train_cl_csv, index = False)
\ No newline at end of file
diff --git a/pymic/net_run_nll/nll_co_teaching.py b/pymic/net_run/noisy_label/nll_co_teaching.py
similarity index 91%
rename from pymic/net_run_nll/nll_co_teaching.py
rename to pymic/net_run/noisy_label/nll_co_teaching.py
index 06240d6..ec8e230 100644
--- a/pymic/net_run_nll/nll_co_teaching.py
+++ b/pymic/net_run/noisy_label/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):
     """
@@ -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()
@@ -147,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):
@@ -162,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)
@@ -174,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']) + "]") 
diff --git a/pymic/net_run_nll/nll_dast.py b/pymic/net_run/noisy_label/nll_dast.py
similarity index 96%
rename from pymic/net_run_nll/nll_dast.py
rename to pymic/net_run/noisy_label/nll_dast.py
index f057f64..1921e9c 100644
--- a/pymic/net_run_nll/nll_dast.py
+++ b/pymic/net_run/noisy_label/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
@@ -156,7 +155,7 @@ def worker_init_fn(worker_id):
                 worker_init = None
 
             bn_train_noise = self.config['dataset']['train_batch_size_noise']
-            num_worker = self.config['dataset'].get('num_workder', 16)
+            num_worker = self.config['dataset'].get('num_worker', 16)
             self.train_loader_noise = torch.utils.data.DataLoader(self.train_set_noise, 
                 batch_size = bn_train_noise, shuffle=True, num_workers= num_worker,
                 worker_init_fn=worker_init)
@@ -239,9 +238,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()
@@ -260,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):
diff --git a/pymic/net_run_nll/nll_trinet.py b/pymic/net_run/noisy_label/nll_trinet.py
similarity index 91%
rename from pymic/net_run_nll/nll_trinet.py
rename to pymic/net_run/noisy_label/nll_trinet.py
index 7ac74af..25c90cf 100644
--- a/pymic/net_run_nll/nll_trinet.py
+++ b/pymic/net_run/noisy_label/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()
@@ -144,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):
@@ -159,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)
@@ -171,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']) + "]") 
diff --git a/pymic/net_run/net_run.py b/pymic/net_run/predict.py
similarity index 51%
rename from pymic/net_run/net_run.py
rename to pymic/net_run/predict.py
index aec6fa0..ca4ef25 100644
--- a/pymic/net_run/net_run.py
+++ b/pymic/net_run/predict.py
@@ -3,6 +3,7 @@
 import logging
 import os
 import sys
+from datetime import datetime
 from pymic.util.parse_config import *
 from pymic.net_run.agent_cls import ClassificationAgent
 from pymic.net_run.agent_seg import SegmentationAgent
@@ -11,27 +12,33 @@ def main():
     """
     The main function for running a network for training or inference.
     """
-    if(len(sys.argv) < 3):
-        print('Number of arguments should be 3. e.g.')
-        print('   pymic_run train config.cfg')
+    if(len(sys.argv) < 2):
+        print('Number of arguments should be 2. e.g.')
+        print('   pymic_test config.cfg')
         exit()
-    stage    = str(sys.argv[1])
-    cfg_file = str(sys.argv[2])
+    cfg_file = str(sys.argv[1])
+    if(not os.path.isfile(cfg_file)):
+        raise ValueError("The config file does not exist: " + cfg_file)
     config   = parse_config(cfg_file)
     config   = synchronize_config(config)
-    log_dir  = config['training']['ckpt_save_dir']
+    log_dir  = config['testing']['output_dir']
     if(not os.path.exists(log_dir)):
         os.makedirs(log_dir, exist_ok=True)
-    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_test.txt", 
+                            level=logging.INFO, format='%(message)s', force=True) # for python 3.9
+    else:
+        logging.basicConfig(filename=log_dir+"/log_test.txt", 
+                            level=logging.INFO, format='%(message)s') # for python 3.6
     logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
     logging_config(config)
     task     = config['dataset']['task_type']
     assert task in ['cls', 'cls_nexcl', 'seg']
     if(task == 'cls' or task == 'cls_nexcl'):
-        agent = ClassificationAgent(config, stage)
+        agent = ClassificationAgent(config, 'test')
     else:
-        agent = SegmentationAgent(config, stage)
+        agent = SegmentationAgent(config, 'test')
     agent.run()
 
 if __name__ == "__main__":
diff --git a/pymic/net_run/self_sup/__init__.py b/pymic/net_run/self_sup/__init__.py
new file mode 100644
index 0000000..55f26bf
--- /dev/null
+++ b/pymic/net_run/self_sup/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from pymic.net_run.self_sup.self_sl_agent import SelfSLSegAgent
\ No newline at end of file
diff --git a/pymic/net_run/self_sup/self_sl_agent.py b/pymic/net_run/self_sup/self_sl_agent.py
new file mode 100644
index 0000000..24a6e66
--- /dev/null
+++ b/pymic/net_run/self_sup/self_sl_agent.py
@@ -0,0 +1,245 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+import copy
+import logging
+import time
+import logging
+import numpy as np
+import random
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from datetime import datetime
+from random import random
+from torch.optim import lr_scheduler
+from tensorboardX import SummaryWriter
+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
+from pymic.loss.seg.util import get_classwise_dice
+from pymic.net_run.infer_func import Inferer
+from pymic.net_run.agent_seg import SegmentationAgent
+from pymic.transform.trans_dict import TransformDict
+from pymic.loss.seg.mse import MAELoss, MSELoss
+
+RegressionLossDict = {
+    'MAELoss': MAELoss,
+    'MSELoss': MSELoss
+    }
+
+class SelfSLSegAgent(SegmentationAgent):
+    """
+    Abstract class for self-supervised segmentation.
+
+    :param config: (dict) A dictionary containing the configuration.
+    :param stage: (str) One of the stage in `train` (default), `inference` or `test`. 
+
+    .. note::
+
+        In the configuration dictionary, in addition to the four sections (`dataset`,
+        `network`, `training` and `inference`) used in fully supervised learning, an 
+        extra section `semi_supervised_learning` is needed. See :doc:`usage.ssl` for details.
+    """
+    def __init__(self, config, stage = 'train'):
+        super(SelfSLSegAgent, self).__init__(config, stage)
+        self.transform_dict  = TransformDict
+
+    def create_loss_calculator(self):
+        if(self.loss_dict is None):
+            self.loss_dict = RegressionLossDict
+        loss_name = self.config['training']['loss_type']
+        if isinstance(loss_name, (list, tuple)):
+            raise ValueError("Undefined loss function {0:}".format(loss_name))
+        elif (loss_name not in self.loss_dict):
+            raise ValueError("Undefined loss function {0:}".format(loss_name))
+        else:
+            loss_param = self.config['training']
+            loss_param['loss_softmax'] = False
+            base_loss = self.loss_dict[loss_name](self.config['training'])
+        if(self.config['training'].get('deep_supervise', False)):
+            raise ValueError("Deep supervised loss not implemented for self-supervised learning")
+            # 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 training(self):
+        iter_valid  = self.config['training']['iter_valid']
+        train_loss  = 0
+        self.net.train()
+        for it in range(iter_valid):
+            try:
+                data = next(self.trainIter)
+            except StopIteration:
+                self.trainIter = iter(self.train_loader)
+                data = next(self.trainIter)
+            # get the inputs
+            inputs  = self.convert_tensor_type(data['image'])
+            label   = self.convert_tensor_type(data['label'])                 
+                   
+            # for debug
+            # from pymic.io.image_read_write import save_nd_array_as_image
+            # for i in range(inputs.shape[0]):
+            #     image_i = inputs[i][0]
+            #     image_name = "temp/image_{0:}_{1:}.nii.gz".format(it, i)
+            #     save_nd_array_as_image(image_i, image_name, reference_name = None)
+            # return
+
+            inputs, label = inputs.to(self.device), label.to(self.device)
+            
+            # zero the parameter gradients
+            self.optimizer.zero_grad()
+                
+            # forward + backward + optimize
+            outputs = self.net(inputs)
+            outputs = nn.Sigmoid()(outputs)
+            loss = self.get_loss_value(data, outputs, label)
+            loss.backward()
+            self.optimizer.step()
+            train_loss = train_loss + loss.item()
+            # get dice evaluation for each class
+            if(isinstance(outputs, tuple) or isinstance(outputs, list)):
+                outputs = outputs[0] 
+
+        train_avg_loss = train_loss / iter_valid
+        train_scalers = {'loss': train_avg_loss}
+        return train_scalers
+        
+    def validation(self):
+        if(self.inferer is None):
+            infer_cfg = self.config['testing']
+            self.inferer = Inferer(infer_cfg)
+        
+        valid_loss_list = []
+        validIter  = iter(self.valid_loader)
+        with torch.no_grad():
+            self.net.eval()
+            for data in validIter:
+                inputs = self.convert_tensor_type(data['image'])
+                label  = self.convert_tensor_type(data['label'])
+                inputs, label  = inputs.to(self.device), label.to(self.device)
+                outputs = self.inferer.run(self.net, inputs)
+                outputs = nn.Sigmoid()(outputs)
+                # The tensors are on CPU when calculating loss for validation data
+                loss = self.get_loss_value(data, outputs, label)
+                valid_loss_list.append(loss.item())
+
+        valid_avg_loss = np.asarray(valid_loss_list).mean()
+        valid_scalers = {'loss': valid_avg_loss}
+        return valid_scalers
+
+    def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
+        loss_scalar ={'train':train_scalars['loss'], 
+                      'valid':valid_scalars['loss']}
+        self.summ_writer.add_scalars('loss', loss_scalar, glob_it)
+        self.summ_writer.add_scalars('lr', {"lr": lr_value}, glob_it)
+        logging.info('train loss {0:.4f}'.format(train_scalars['loss']))        
+        logging.info('valid loss {0:.4f}'.format(valid_scalars['loss']))  
+
+    def train_valid(self):
+        device_ids = self.config['training']['gpus']
+        if(len(device_ids) > 1):
+            self.device = torch.device("cuda:0")
+            self.net = nn.DataParallel(self.net, device_ids = device_ids)
+        else:
+            self.device = torch.device("cuda:{0:}".format(device_ids[0]))
+        self.net.to(self.device)
+        ckpt_dir    = self.config['training']['ckpt_save_dir']
+        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_max    = self.config['training']['iter_max']
+        iter_valid  = self.config['training']['iter_valid']
+        iter_save   = self.config['training'].get('iter_save', None)
+        early_stop_it = self.config['training'].get('early_stop_patience', None)
+        if(iter_save is None):
+            iter_save_list = [iter_max]
+        elif(isinstance(iter_save, (tuple, list))):
+            iter_save_list = iter_save
+        else:
+            iter_save_list = range(0, iter_max + 1, iter_save)
+
+        self.min_val_loss = 10000.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'])
+            else:
+                self.net.load_state_dict(self.checkpoint['model_state_dict'])
+            self.min_val_loss = self.checkpoint.get('valid_loss', 10000)
+            # 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.create_loss_calculator()
+    
+        self.trainIter  = iter(self.train_loader)
+        
+        logging.info("{0:} training start".format(str(datetime.now())[:-7]))
+        self.summ_writer = SummaryWriter(self.config['training']['ckpt_save_dir'])
+        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['loss'])
+            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['loss'] < self.min_val_loss):
+                self.min_val_loss = valid_scalars['loss']
+                self.max_val_it   = self.glob_it
+                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
+            if ((self.glob_it in iter_save_list) or stop_now):
+                save_dict = {'iteration': self.glob_it,
+                             'valid_loss': valid_scalars['loss'],
+                             'model_state_dict': self.net.module.state_dict() \
+                                 if len(device_ids) > 1 else self.net.state_dict(),
+                             'optimizer_state_dict': self.optimizer.state_dict()}
+                save_name = "{0:}/{1:}_{2:}.pt".format(ckpt_dir, ckpt_prefix, self.glob_it)
+                torch.save(save_dict, save_name) 
+                txt_file = open("{0:}/{1:}_latest.txt".format(ckpt_dir, ckpt_prefix), 'wt')
+                txt_file.write(str(self.glob_it))
+                txt_file.close()
+            if(stop_now):
+                logging.info("The training is early stopped")
+                break
+        # save the best performing checkpoint
+        save_dict = {'iteration': self.max_val_it,
+                    'valid_loss': self.min_val_loss,
+                    'model_state_dict': self.best_model_wts,
+                    'optimizer_state_dict': self.optimizer.state_dict()}
+        save_name = "{0:}/{1:}_{2:}.pt".format(ckpt_dir, ckpt_prefix, self.max_val_it)
+        torch.save(save_dict, save_name) 
+        txt_file = open("{0:}/{1:}_best.txt".format(ckpt_dir, ckpt_prefix), 'wt')
+        txt_file.write(str(self.max_val_it))
+        txt_file.close()
+        logging.info('The best performing iter is {0:}, valid loss {1:}'.format(\
+            self.max_val_it, self.min_val_loss))
+        self.summ_writer.close()
\ No newline at end of file
diff --git a/pymic/net_run/semi_sup/__init__.py b/pymic/net_run/semi_sup/__init__.py
new file mode 100644
index 0000000..be753c2
--- /dev/null
+++ b/pymic/net_run/semi_sup/__init__.py
@@ -0,0 +1,16 @@
+from __future__ import absolute_import
+from pymic.net_run.semi_sup.ssl_abstract import SSLSegAgent
+from pymic.net_run.semi_sup.ssl_em import SSLEntropyMinimization
+from pymic.net_run.semi_sup.ssl_mt import SSLMeanTeacher
+from pymic.net_run.semi_sup.ssl_uamt import SSLUncertaintyAwareMeanTeacher
+from pymic.net_run.semi_sup.ssl_cct import SSLCCT
+from pymic.net_run.semi_sup.ssl_cps import SSLCPS
+from pymic.net_run.semi_sup.ssl_urpc import SSLURPC
+
+
+SSLMethodDict = {'EntropyMinimization': SSLEntropyMinimization,
+    'MeanTeacher': SSLMeanTeacher,
+    'UAMT': SSLUncertaintyAwareMeanTeacher,
+    'CCT': SSLCCT,
+    'CPS': SSLCPS,
+    'URPC': SSLURPC}
\ No newline at end of file
diff --git a/pymic/net_run_ssl/ssl_abstract.py b/pymic/net_run/semi_sup/ssl_abstract.py
similarity index 90%
rename from pymic/net_run_ssl/ssl_abstract.py
rename to pymic/net_run/semi_sup/ssl_abstract.py
index e847d6c..5a46257 100644
--- a/pymic/net_run_ssl/ssl_abstract.py
+++ b/pymic/net_run/semi_sup/ssl_abstract.py
@@ -73,7 +73,7 @@ def worker_init_fn(worker_id):
                 worker_init = None
 
             bn_train_unlab = self.config['dataset']['train_batch_size_unlab']
-            num_worker = self.config['dataset'].get('num_workder', 16)
+            num_worker = self.config['dataset'].get('num_worker', 16)
             self.train_loader_unlab = torch.utils.data.DataLoader(self.train_set_unlab, 
                 batch_size = bn_train_unlab, shuffle=True, num_workers= num_worker,
                 worker_init_fn=worker_init)
@@ -83,7 +83,7 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
                       'valid':valid_scalars['loss']}
         loss_sup_scalar  = {'train':train_scalars['loss_sup']}
         loss_upsup_scalar  = {'train':train_scalars['loss_reg']}
-        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_sup', loss_sup_scalar, glob_it)
         self.summ_writer.add_scalars('loss_reg', loss_upsup_scalar, glob_it)
@@ -95,11 +95,11 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
             cls_dice_scalar = {'train':train_scalars['class_dice'][c], \
                 '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']) + "]")  
 
     def train_valid(self):
diff --git a/pymic/net_run_ssl/ssl_cct.py b/pymic/net_run/semi_sup/ssl_cct.py
similarity index 94%
rename from pymic/net_run_ssl/ssl_cct.py
rename to pymic/net_run/semi_sup/ssl_cct.py
index 0bc23a2..1943608 100644
--- a/pymic/net_run_ssl/ssl_cct.py
+++ b/pymic/net_run/semi_sup/ssl_cct.py
@@ -5,11 +5,10 @@
 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
-from pymic.net_run_ssl.ssl_abstract import SSLSegAgent
+from pymic.net_run.semi_sup import SSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 
 def softmax_mse_loss(inputs, targets, conf_mask=False, threshold=None, use_softmax=False):
@@ -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() 
@@ -157,9 +154,9 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
diff --git a/pymic/net_run_ssl/ssl_cps.py b/pymic/net_run/semi_sup/ssl_cps.py
similarity index 92%
rename from pymic/net_run_ssl/ssl_cps.py
rename to pymic/net_run/semi_sup/ssl_cps.py
index 239443c..4a3be9c 100644
--- a/pymic/net_run_ssl/ssl_cps.py
+++ b/pymic/net_run/semi_sup/ssl_cps.py
@@ -4,11 +4,10 @@
 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
-from pymic.net_run_ssl.ssl_abstract import SSLSegAgent
+from pymic.net_run.semi_sup import SSLSegAgent
 from pymic.net.net_dict_seg import SegNetDict
 from pymic.util.ramps import get_rampup_ratio
 
@@ -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()
@@ -141,12 +137,12 @@ def training(self):
         train_avg_loss_pse_sup1 = train_loss_pseudo_sup1 / iter_valid 
         train_avg_loss_pse_sup2 = train_loss_pseudo_sup2 / 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_sup1':train_avg_loss_sup1, 'loss_sup2': train_avg_loss_sup2,
             'loss_pse_sup1':train_avg_loss_pse_sup1, 'loss_pse_sup2': train_avg_loss_pse_sup2,
-            'regular_w':regular_w, 'avg_dice':train_avg_dice, 'class_dice': train_cls_dice}
+            'regular_w':regular_w, '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):
@@ -156,7 +152,7 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
                             'net2':train_scalars['loss_sup2']}
         loss_pse_sup_scalar = {'net1':train_scalars['loss_pse_sup1'],
                                'net2':train_scalars['loss_pse_sup2']}
-        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_sup', loss_sup_scalar, glob_it)
         self.summ_writer.add_scalars('loss_pseudo_sup', loss_pse_sup_scalar, glob_it)
@@ -170,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']) + "]") 
\ No newline at end of file
diff --git a/pymic/net_run_ssl/ssl_em.py b/pymic/net_run/semi_sup/ssl_em.py
similarity index 92%
rename from pymic/net_run_ssl/ssl_em.py
rename to pymic/net_run/semi_sup/ssl_em.py
index 745501f..fde941b 100644
--- a/pymic/net_run_ssl/ssl_em.py
+++ b/pymic/net_run/semi_sup/ssl_em.py
@@ -3,12 +3,11 @@
 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
 from pymic.loss.seg.ssl import EntropyLoss
-from pymic.net_run_ssl.ssl_abstract import SSLSegAgent
+from pymic.net_run.semi_sup import SSLSegAgent
 from pymic.transform.trans_dict import TransformDict
 from pymic.util.ramps import get_rampup_ratio
 
@@ -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()
@@ -102,9 +98,9 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
\ No newline at end of file
diff --git a/pymic/net_run_ssl/ssl_mt.py b/pymic/net_run/semi_sup/ssl_mt.py
similarity index 92%
rename from pymic/net_run_ssl/ssl_mt.py
rename to pymic/net_run/semi_sup/ssl_mt.py
index bf4dacd..2a2abb8 100644
--- a/pymic/net_run_ssl/ssl_mt.py
+++ b/pymic/net_run/semi_sup/ssl_mt.py
@@ -3,11 +3,10 @@
 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
-from pymic.net_run_ssl.ssl_abstract import SSLSegAgent
+from pymic.net_run.semi_sup import SSLSegAgent
 from pymic.net.net_dict_seg import SegNetDict
 from pymic.util.ramps import get_rampup_ratio
 
@@ -102,13 +101,10 @@ 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)
-            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)
 
@@ -127,9 +123,9 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
\ No newline at end of file
diff --git a/pymic/net_run_ssl/ssl_uamt.py b/pymic/net_run/semi_sup/ssl_uamt.py
similarity index 93%
rename from pymic/net_run_ssl/ssl_uamt.py
rename to pymic/net_run/semi_sup/ssl_uamt.py
index 40a742b..6222fe3 100644
--- a/pymic/net_run_ssl/ssl_uamt.py
+++ b/pymic/net_run/semi_sup/ssl_uamt.py
@@ -3,11 +3,10 @@
 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
-from pymic.net_run_ssl.ssl_mt import SSLMeanTeacher
+from pymic.net_run.semi_sup import SSLMeanTeacher
 from pymic.util.ramps import get_rampup_ratio
 
 class SSLUncertaintyAwareMeanTeacher(SSLMeanTeacher):
@@ -104,14 +103,10 @@ 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)
-            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)
 
@@ -130,9 +125,9 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
\ No newline at end of file
diff --git a/pymic/net_run_ssl/ssl_urpc.py b/pymic/net_run/semi_sup/ssl_urpc.py
similarity index 92%
rename from pymic/net_run_ssl/ssl_urpc.py
rename to pymic/net_run/semi_sup/ssl_urpc.py
index 336e525..56bb77e 100644
--- a/pymic/net_run_ssl/ssl_urpc.py
+++ b/pymic/net_run/semi_sup/ssl_urpc.py
@@ -4,11 +4,10 @@
 import torch
 import torch.nn as nn
 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
-from pymic.net_run_ssl.ssl_abstract import SSLSegAgent
+from pymic.net_run.semi_sup import SSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 
 class SSLURPC(SSLSegAgent):
@@ -72,8 +71,11 @@ def training(self):
             p0 = [output_i[:n0] for output_i in outputs_list]
             loss_sup = self.get_loss_value(data_lab, p0, y0)
 
-            # get average probability across scales
+            # resize to the same shape, and get average probability across scales
             outputs_soft_list = [torch.softmax(item, dim=1) for item in outputs_list]
+            for i in range(1, len(outputs_soft_list)):
+                outputs_soft_list[i] = nn.functional.interpolate(outputs_soft_list[i],
+                    outputs_soft_list[0].shape[2:])
             outputs_soft_avg  = torch.mean(torch.stack(outputs_soft_list),dim = 0)
             p1_avg = outputs_soft_avg[n0:] * 0.99 + 0.005 # for unannotated images
 
@@ -96,9 +98,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() 
@@ -114,9 +114,9 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,  'class_dice': train_cls_dice}
         return train_scalers
diff --git a/pymic/net_run/train.py b/pymic/net_run/train.py
new file mode 100644
index 0000000..1478527
--- /dev/null
+++ b/pymic/net_run/train.py
@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+import logging
+import os
+import sys
+import shutil
+from datetime import datetime
+from pymic.util.parse_config import *
+from pymic.net_run.agent_cls import ClassificationAgent
+from pymic.net_run.agent_seg import SegmentationAgent
+from pymic.net_run.semi_sup import SSLMethodDict
+from pymic.net_run.weak_sup import WSLMethodDict
+from pymic.net_run.noisy_label import NLLMethodDict
+from pymic.net_run.self_sup import SelfSLSegAgent
+
+def get_segmentation_agent(config, sup_type):
+    assert(sup_type in ['fully_sup', 'semi_sup', 'self_sup', 'weak_sup', 'noisy_label'])
+    if(sup_type == 'fully_sup'):
+        logging.info("\n********** Fully Supervised Learning **********\n")
+        agent = SegmentationAgent(config, 'train')
+    elif(sup_type == 'semi_sup'):
+        logging.info("\n********** Semi Supervised Learning **********\n")
+        method = config['semi_supervised_learning']['method_name']
+        agent = SSLMethodDict[method](config, 'train')
+    elif(sup_type == 'weak_sup'):
+        logging.info("\n********** Weakly Supervised Learning **********\n")
+        method = config['weakly_supervised_learning']['method_name']
+        agent = WSLMethodDict[method](config, 'train')
+    elif(sup_type == 'noisy_label'):
+        logging.info("\n********** Noisy Label Learning **********\n")
+        method = config['noisy_label_learning']['method_name']
+        agent = NLLMethodDict[method](config, 'train')
+    elif(sup_type == 'self_sup'):
+        logging.info("\n********** Self Supervised Learning **********\n")
+        method = config['self_supervised_learning']['method_name']
+        if(method == "custom"):
+            pass
+        elif(method == "model_genesis"):
+            transforms = ['RandomFlip', 'LocalShuffling', 'NonLinearTransform', 'InOutPainting']
+            genesis_cfg = {
+                'randomflip_flip_depth': True,
+                'randomflip_flip_height': True,
+                'randomflip_flip_width': True,
+                'localshuffling_probability': 0.5,
+                'nonLineartransform_probability': 0.9,
+                'inoutpainting_probability': 0.9,
+                'inpainting_probability': 0.2
+            }
+            config['dataset']['train_transform'].extend(transforms)
+            config['dataset']['valid_transform'].extend(transforms)
+            config['dataset'].update(genesis_cfg)
+            logging_config(config['dataset'])
+        else:
+            raise ValueError("The specified method {0:} is not implemented. ".format(method) + \
+                         "Consider to set `self_sl_method = custom` and use customized" + \
+                         " transforms for self-supervised learning.")
+        agent = SelfSLSegAgent(config, 'train')
+    else:
+        raise ValueError("undefined supervision type: {0:}".format(sup_type))
+    return agent
+
+def main():
+    """
+    The main function for running a network for training.
+    """
+    if(len(sys.argv) < 2):
+        print('Number of arguments should be 2. e.g.')
+        print('   pymic_train config.cfg')
+        exit()
+    cfg_file = str(sys.argv[1])
+    if(not os.path.isfile(cfg_file)):
+        raise ValueError("The config file does not exist: " + cfg_file)
+    config   = parse_config(cfg_file)
+    config   = synchronize_config(config)
+    log_dir  = config['training']['ckpt_save_dir']
+    if(not os.path.exists(log_dir)):
+        os.makedirs(log_dir, exist_ok=True)
+    dst_cfg = cfg_file if "/" not in cfg_file else cfg_file.split("/")[-1]
+    shutil.copy(cfg_file, log_dir + "/" + dst_cfg)
+    if sys.version.startswith("3.9"):
+        logging.basicConfig(filename=log_dir+"/log_train_{0:}.txt".format(str(datetime.now())[:-7]), 
+                            level=logging.INFO, format='%(message)s', force=True) # for python 3.9
+    else:
+        logging.basicConfig(filename=log_dir+"/log_train_{0:}.txt".format(str(datetime.now())[:-7]), 
+                            level=logging.INFO, format='%(message)s') # for python 3.6
+    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
+    logging_config(config)
+    task     = config['dataset']['task_type']
+    assert task in ['cls', 'cls_nexcl', 'seg']
+    if(task == 'cls' or task == 'cls_nexcl'):
+        agent = ClassificationAgent(config, 'train')
+    else:
+        sup_type = config['dataset'].get('supervise_type', 'fully_sup')
+        agent = get_segmentation_agent(config, sup_type)
+    agent.run()
+
+if __name__ == "__main__":
+    main()
+    
+
diff --git a/pymic/net_run/weak_sup/__init__.py b/pymic/net_run/weak_sup/__init__.py
new file mode 100644
index 0000000..b3c8332
--- /dev/null
+++ b/pymic/net_run/weak_sup/__init__.py
@@ -0,0 +1,15 @@
+from __future__ import absolute_import
+from pymic.net_run.weak_sup.wsl_abstract import WSLSegAgent
+from pymic.net_run.weak_sup.wsl_em import WSLEntropyMinimization
+from pymic.net_run.weak_sup.wsl_gatedcrf import WSLGatedCRF
+from pymic.net_run.weak_sup.wsl_mumford_shah import WSLMumfordShah
+from pymic.net_run.weak_sup.wsl_tv import WSLTotalVariation
+from pymic.net_run.weak_sup.wsl_ustm import WSLUSTM
+from pymic.net_run.weak_sup.wsl_dmpls import WSLDMPLS
+
+WSLMethodDict = {'EntropyMinimization': WSLEntropyMinimization,
+    'GatedCRF': WSLGatedCRF,
+    'MumfordShah': WSLMumfordShah,
+    'TotalVariation': WSLTotalVariation,
+    'USTM': WSLUSTM,
+    'DMPLS': WSLDMPLS}
\ No newline at end of file
diff --git a/pymic/net_run_wsl/wsl_abstract.py b/pymic/net_run/weak_sup/wsl_abstract.py
similarity index 83%
rename from pymic/net_run_wsl/wsl_abstract.py
rename to pymic/net_run/weak_sup/wsl_abstract.py
index f317dd0..f290465 100644
--- a/pymic/net_run_wsl/wsl_abstract.py
+++ b/pymic/net_run/weak_sup/wsl_abstract.py
@@ -24,7 +24,7 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
                       'valid':valid_scalars['loss']}
         loss_sup_scalar  = {'train':train_scalars['loss_sup']}
         loss_upsup_scalar  = {'train':train_scalars['loss_reg']}
-        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_sup', loss_sup_scalar, glob_it)
         self.summ_writer.add_scalars('loss_reg', loss_upsup_scalar, glob_it)
@@ -36,9 +36,9 @@ def write_scalars(self, train_scalars, valid_scalars, lr_value, glob_it):
             cls_dice_scalar = {'train':train_scalars['class_dice'][c], \
                 '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']) + "]")  
diff --git a/pymic/net_run_wsl/wsl_dmpls.py b/pymic/net_run/weak_sup/wsl_dmpls.py
similarity index 93%
rename from pymic/net_run_wsl/wsl_dmpls.py
rename to pymic/net_run/weak_sup/wsl_dmpls.py
index 15e74a0..4212409 100644
--- a/pymic/net_run_wsl/wsl_dmpls.py
+++ b/pymic/net_run/weak_sup/wsl_dmpls.py
@@ -9,7 +9,7 @@
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
 from pymic.loss.seg.dice import DiceLoss
-from pymic.net_run_wsl.wsl_abstract import WSLSegAgent
+from pymic.net_run.weak_sup import WSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 
 class WSLDMPLS(WSLSegAgent):
@@ -91,10 +91,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() 
@@ -110,10 +107,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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+
         return train_scalers
         
\ No newline at end of file
diff --git a/pymic/net_run_wsl/wsl_em.py b/pymic/net_run/weak_sup/wsl_em.py
similarity index 91%
rename from pymic/net_run_wsl/wsl_em.py
rename to pymic/net_run/weak_sup/wsl_em.py
index 387bd93..adcd70c 100644
--- a/pymic/net_run_wsl/wsl_em.py
+++ b/pymic/net_run/weak_sup/wsl_em.py
@@ -3,13 +3,12 @@
 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
 from pymic.loss.seg.ssl import EntropyLoss
 from pymic.net_run.agent_seg import SegmentationAgent
-from pymic.net_run_wsl.wsl_abstract import WSLSegAgent
+from pymic.net_run.weak_sup import WSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 
 class WSLEntropyMinimization(WSLSegAgent):
@@ -72,9 +71,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()
@@ -91,9 +87,9 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
\ No newline at end of file
diff --git a/pymic/net_run_wsl/wsl_gatedcrf.py b/pymic/net_run/weak_sup/wsl_gatedcrf.py
similarity index 93%
rename from pymic/net_run_wsl/wsl_gatedcrf.py
rename to pymic/net_run/weak_sup/wsl_gatedcrf.py
index 85f601e..2ce1f95 100644
--- a/pymic/net_run_wsl/wsl_gatedcrf.py
+++ b/pymic/net_run/weak_sup/wsl_gatedcrf.py
@@ -3,12 +3,11 @@
 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
 from pymic.loss.seg.gatedcrf import GatedCRFLoss
-from pymic.net_run_wsl.wsl_abstract import WSLSegAgent
+from pymic.net_run.weak_sup import WSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 
 class WSLGatedCRF(WSLSegAgent):
@@ -98,9 +97,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()
@@ -117,10 +113,10 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
         
\ No newline at end of file
diff --git a/pymic/net_run_wsl/wsl_mumford_shah.py b/pymic/net_run/weak_sup/wsl_mumford_shah.py
similarity index 91%
rename from pymic/net_run_wsl/wsl_mumford_shah.py
rename to pymic/net_run/weak_sup/wsl_mumford_shah.py
index 862d761..2480fee 100644
--- a/pymic/net_run_wsl/wsl_mumford_shah.py
+++ b/pymic/net_run/weak_sup/wsl_mumford_shah.py
@@ -3,12 +3,11 @@
 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
 from pymic.loss.seg.mumford_shah import MumfordShahLoss
-from pymic.net_run_wsl.wsl_abstract import WSLSegAgent
+from pymic.net_run.weak_sup import WSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 
 class WSLMumfordShah(WSLSegAgent):
@@ -73,9 +72,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()
@@ -92,10 +88,10 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
         
\ No newline at end of file
diff --git a/pymic/net_run_wsl/wsl_tv.py b/pymic/net_run/weak_sup/wsl_tv.py
similarity index 91%
rename from pymic/net_run_wsl/wsl_tv.py
rename to pymic/net_run/weak_sup/wsl_tv.py
index 3fd55a3..9d13c5d 100644
--- a/pymic/net_run_wsl/wsl_tv.py
+++ b/pymic/net_run/weak_sup/wsl_tv.py
@@ -3,12 +3,11 @@
 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
 from pymic.loss.seg.ssl import TotalVariationLoss
-from pymic.net_run_wsl.wsl_abstract import WSLSegAgent
+from pymic.net_run.weak_sup import WSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 from pymic.util.general import keyword_match
 
@@ -68,9 +67,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()
@@ -87,10 +83,10 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
         
\ No newline at end of file
diff --git a/pymic/net_run_wsl/wsl_ustm.py b/pymic/net_run/weak_sup/wsl_ustm.py
similarity index 93%
rename from pymic/net_run_wsl/wsl_ustm.py
rename to pymic/net_run/weak_sup/wsl_ustm.py
index b033f9d..0ea3fbc 100644
--- a/pymic/net_run_wsl/wsl_ustm.py
+++ b/pymic/net_run/weak_sup/wsl_ustm.py
@@ -5,12 +5,11 @@
 import random
 import torch
 import torch.nn.functional as F
-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
 from pymic.net.net_dict_seg import SegNetDict
-from pymic.net_run_wsl.wsl_abstract import WSLSegAgent
+from pymic.net_run.weak_sup import WSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 from pymic.util.general import keyword_match
 
@@ -121,13 +120,10 @@ 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 = wsl_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)
 
@@ -146,9 +142,9 @@ 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':regular_w,
-            'avg_dice':train_avg_dice,     'class_dice': train_cls_dice}
+            'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice}
         return train_scalers
\ No newline at end of file
diff --git a/pymic/net_run_nll/__init__.py b/pymic/net_run_nll/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/pymic/net_run_nll/nll_main.py b/pymic/net_run_nll/nll_main.py
deleted file mode 100644
index a855f2b..0000000
--- a/pymic/net_run_nll/nll_main.py
+++ /dev/null
@@ -1,42 +0,0 @@
-
-# -*- coding: utf-8 -*-
-from __future__ import print_function, division
-import logging 
-import os
-import sys
-from pymic.util.parse_config import *
-from pymic.net_run_nll.nll_co_teaching import NLLCoTeaching
-from pymic.net_run_nll.nll_trinet import NLLTriNet
-from pymic.net_run_nll.nll_dast import NLLDAST
-
-NLLMethodDict = {'CoTeaching': NLLCoTeaching,
-    "TriNet": NLLTriNet,
-    "DAST": NLLDAST}
-
-def main():
-    """
-    The main function for noisy label learning methods.
-    """
-    if(len(sys.argv) < 3):
-        print('Number of arguments should be 3. e.g.')
-        print('   pymic_nll train config.cfg')
-        exit()
-    stage    = str(sys.argv[1])
-    cfg_file = str(sys.argv[2])
-    config   = parse_config(cfg_file)
-    config   = synchronize_config(config)
-    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')
-    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
-    logging_config(config)
-    nll_method = config['noisy_label_learning']['nll_method']
-    agent = NLLMethodDict[nll_method](config, stage)
-    agent.run()
-
-if __name__ == "__main__":
-    main()
-
-    
\ No newline at end of file
diff --git a/pymic/net_run_ssl/__init__.py b/pymic/net_run_ssl/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/pymic/net_run_ssl/ssl_main.py b/pymic/net_run_ssl/ssl_main.py
deleted file mode 100644
index 7c4f2b9..0000000
--- a/pymic/net_run_ssl/ssl_main.py
+++ /dev/null
@@ -1,49 +0,0 @@
-
-# -*- coding: utf-8 -*-
-from __future__ import print_function, division
-import logging 
-import os
-import sys
-from pymic.util.parse_config import *
-from pymic.net_run_ssl.ssl_em import SSLEntropyMinimization
-from pymic.net_run_ssl.ssl_mt import SSLMeanTeacher
-from pymic.net_run_ssl.ssl_uamt import SSLUncertaintyAwareMeanTeacher
-from pymic.net_run_ssl.ssl_cct import SSLCCT
-from pymic.net_run_ssl.ssl_cps import SSLCPS
-from pymic.net_run_ssl.ssl_urpc import SSLURPC
-
-
-SSLMethodDict = {'EntropyMinimization': SSLEntropyMinimization,
-    'MeanTeacher': SSLMeanTeacher,
-    'UAMT': SSLUncertaintyAwareMeanTeacher,
-    'CCT': SSLCCT,
-    'CPS': SSLCPS,
-    'URPC': SSLURPC}
-
-def main():
-    """
-    Main function for running a semi-supervised method.
-    """
-    if(len(sys.argv) < 3):
-        print('Number of arguments should be 3. e.g.')
-        print('   pymic_ssl train config.cfg')
-        exit()
-    stage    = str(sys.argv[1])
-    cfg_file = str(sys.argv[2])
-    config   = parse_config(cfg_file)
-    config   = synchronize_config(config)
-    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')
-    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
-    logging_config(config)
-    ssl_method = config['semi_supervised_learning']['ssl_method']
-    agent = SSLMethodDict[ssl_method](config, stage)
-    agent.run()
-
-if __name__ == "__main__":
-    main()
-
-    
\ No newline at end of file
diff --git a/pymic/net_run_wsl/__init__.py b/pymic/net_run_wsl/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/pymic/net_run_wsl/wsl_main.py b/pymic/net_run_wsl/wsl_main.py
deleted file mode 100644
index d8e791d..0000000
--- a/pymic/net_run_wsl/wsl_main.py
+++ /dev/null
@@ -1,48 +0,0 @@
-
-# -*- coding: utf-8 -*-
-from __future__ import print_function, division
-import logging 
-import os
-import sys
-from pymic.util.parse_config import *
-from pymic.net_run_wsl.wsl_em import WSLEntropyMinimization
-from pymic.net_run_wsl.wsl_gatedcrf import WSLGatedCRF
-from pymic.net_run_wsl.wsl_mumford_shah import WSLMumfordShah
-from pymic.net_run_wsl.wsl_tv import WSLTotalVariation
-from pymic.net_run_wsl.wsl_ustm import WSLUSTM
-from pymic.net_run_wsl.wsl_dmpls import WSLDMPLS
-
-WSLMethodDict = {'EntropyMinimization': WSLEntropyMinimization,
-    'GatedCRF': WSLGatedCRF,
-    'MumfordShah': WSLMumfordShah,
-    'TotalVariation': WSLTotalVariation,
-    'USTM': WSLUSTM,
-    'DMPLS': WSLDMPLS}
-
-def main():
-    """
-    The main function for training and inference of weakly supervised segmentation. 
-    """
-    if(len(sys.argv) < 3):
-        print('Number of arguments should be 3. e.g.')
-        print('   pymic_wsl train config.cfg')
-        exit()
-    stage    = str(sys.argv[1])
-    cfg_file = str(sys.argv[2])
-    config   = parse_config(cfg_file)
-    config   = synchronize_config(config)
-    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')
-    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
-    logging_config(config)
-    wsl_method = config['weakly_supervised_learning']['wsl_method']
-    agent = WSLMethodDict[wsl_method](config, stage)
-    agent.run()
-
-if __name__ == "__main__":
-    main()
-
-    
\ No newline at end of file
diff --git a/pymic/transform/__init__.py b/pymic/transform/__init__.py
index e69de29..b47e09d 100644
--- a/pymic/transform/__init__.py
+++ b/pymic/transform/__init__.py
@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_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
\ No newline at end of file
diff --git a/pymic/transform/intensity.py b/pymic/transform/intensity.py
index 171604b..3b5ee9d 100644
--- a/pymic/transform/intensity.py
+++ b/pymic/transform/intensity.py
@@ -1,15 +1,48 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division
-
-import torch
+import copy 
 import json
 import math
 import random
 import numpy as np
-from scipy import ndimage
 from pymic.transform.abstract_transform import AbstractTransform
 from pymic.util.image_process import *
+try:  # SciPy >= 0.19
+    from scipy.special import comb
+except ImportError:
+    from scipy.misc import comb
+
+def bernstein_poly(i, n, t):
+    """
+     The Bernstein polynomial of n, i as a function of t
+    """
+
+    return comb(n, i) * ( t**(n-i) ) * (1 - t)**i
+
+def bezier_curve(points, nTimes=1000):
+    """
+       Given a set of control points, return the
+       bezier curve defined by the control points.
+       Control points should be a list of lists, or list of tuples
+       such as [ [1,1], 
+                 [2,3], 
+                 [4,5], ..[Xn, Yn] ]
+        nTimes is the number of time steps, defaults to 1000
+        See http://processingjs.nihongoresources.com/bezierinfo/
+    """
 
+    nPoints = len(points)
+    xPoints = np.array([p[0] for p in points])
+    yPoints = np.array([p[1] for p in points])
+
+    t = np.linspace(0.0, 1.0, nTimes)
+
+    polynomial_array = np.array([ bernstein_poly(i, nPoints-1, t) for i in range(0, nPoints)   ])
+    
+    xvals = np.dot(xPoints, polynomial_array)
+    yvals = np.dot(yPoints, polynomial_array)
+
+    return xvals, yvals
 
 class GammaCorrection(AbstractTransform):
     """
@@ -98,4 +131,230 @@ def __call__(self, sample):
         assert(image.shape[0] == 1 or image.shape[0] == 3)
         if(image.shape[0] == 1):
             sample['image'] = np.concatenate([image, image, image])
+        return sample
+    
+class NonLinearTransform(AbstractTransform):
+    def __init__(self, params):
+        super(NonLinearTransform, self).__init__(params)
+        self.inverse  = params.get('NonLinearTransform_inverse'.lower(), False)
+        self.prob     = params.get('NonLinearTransform_probability'.lower(), 0.5)
+    
+    def __call__(self, sample):
+        if(random.random() >  self.prob):
+            return sample
+
+        image= sample['image'] 
+        points = [[0, 0], [random.random(), random.random()], [random.random(), random.random()], [1, 1]]
+        xvals, yvals = bezier_curve(points, nTimes=100000)
+        if random.random() < 0.5: # Half change to get flip
+            xvals = np.sort(xvals)
+        else:
+            xvals, yvals = np.sort(xvals), np.sort(yvals)
+        image = np.interp(image, xvals, yvals)
+        sample['image']  = image 
+        return sample
+
+class LocalShuffling(AbstractTransform):
+    """
+    local pixel shuffling of an input image, used for self-supervised learning 
+    """
+    def __init__(self, params):
+        super(LocalShuffling, self).__init__(params)
+        self.inverse  = params.get('LocalShuffling_inverse'.lower(), False)
+        self.prob     = params.get('LocalShuffling_probability'.lower(), 0.5)
+        self.block_range = params.get('LocalShuffling_block_range'.lower(), (5000, 10000))
+        self.block_size_min = params.get('LocalShuffling_block_size_min'.lower(), None)
+        self.block_size_max = params.get('LocalShuffling_block_size_max'.lower(), None)
+
+    def __call__(self, sample):
+        if(random.random() >  self.prob):
+            return sample
+
+        image= sample['image']       
+        img_shape = image.shape
+        img_dim = len(img_shape) - 1
+        assert(img_dim == 2 or img_dim == 3)
+        img_out = copy.deepcopy(image)
+        if(self.block_size_min is None):
+            block_size_min = [2] * img_dim
+        elif(isinstance(self.block_size_min, int)):
+            block_size_min = [self.block_size_min] * img_dim
+        else:
+            assert(len(self.block_size_min) == img_dim)
+            block_size_min = self.block_size_min
+
+        if(self.block_size_max is None):
+            block_size_max = [img_shape[1+i]//10 for i in range(img_dim)]
+        elif(isinstance(self.block_size_min, int)):
+            block_size_max = [self.block_size_max] * img_dim
+        else:
+            assert(len(self.block_size_max) == img_dim)
+            block_size_max = self.block_size_max
+        block_num = random.randint(self.block_range[0], self.block_range[1])
+
+        for n in range(block_num):
+            block_size = [random.randint(block_size_min[i], block_size_max[i]) \
+                for i in range(img_dim)]    
+            coord_min = [random.randint(0, img_shape[1+i] - block_size[i]) \
+                for i in range(img_dim)]
+            if(img_dim == 2):
+                window = image[:, coord_min[0]:coord_min[0] + block_size[0], 
+                                  coord_min[1]:coord_min[1] + block_size[1]]
+                n_pixels = block_size[0] * block_size[1]
+            else:
+                window = image[:, coord_min[0]:coord_min[0] + block_size[0], 
+                                  coord_min[1]:coord_min[1] + block_size[1],
+                                  coord_min[2]:coord_min[2] + block_size[2]]
+                n_pixels = block_size[0] * block_size[1] * block_size[2]
+            window = np.reshape(window, [-1, n_pixels])
+            np.random.shuffle(np.transpose(window))
+            window = np.transpose(window)
+            if(img_dim == 2):
+                window = np.reshape(window, [-1, block_size[0], block_size[1]])
+                img_out[:, coord_min[0]:coord_min[0] + block_size[0], 
+                           coord_min[1]:coord_min[1] + block_size[1]] = window
+            else:
+                window = np.reshape(window, [-1, block_size[0], block_size[1], block_size[2]])
+                img_out[:, coord_min[0]:coord_min[0] + block_size[0], 
+                           coord_min[1]:coord_min[1] + block_size[1],
+                           coord_min[2]:coord_min[2] + block_size[2]] = window
+        sample['image'] = img_out
+        return sample
+
+class InPainting(AbstractTransform):
+    """
+    In-painting of an input image, used for self-supervised learning 
+    """
+    def __init__(self, params):
+        super(InPainting, self).__init__(params)
+        self.inverse  = params.get('InPainting_inverse'.lower(), False)
+        self.prob     = params.get('InPainting_probability'.lower(), 0.5)
+        self.block_range = params.get('InPainting_block_range'.lower(), (1, 6))
+        self.block_size_min = params.get('InPainting_block_size_min'.lower(), None)
+        self.block_size_max = params.get('InPainting_block_size_max'.lower(), None)
+
+    def __call__(self, sample):
+        if(random.random() >  self.prob):
+            return sample
+
+        image= sample['image']       
+        img_shape = image.shape
+        img_dim = len(img_shape) - 1
+        assert(img_dim == 2 or img_dim == 3)
+
+        if(self.block_size_min is None):
+            block_size_min = [img_shape[1+i]//6 for i in range(img_dim)]
+        elif(isinstance(self.block_size_min, int)):
+            block_size_min = [self.block_size_min] * img_dim
+        else:
+            assert(len(self.block_size_min) == img_dim)
+            block_size_min = self.block_size_min
+
+        if(self.block_size_max is None):
+            block_size_max = [img_shape[1+i]//3 for i in range(img_dim)]
+        elif(isinstance(self.block_size_min, int)):
+            block_size_max = [self.block_size_max] * img_dim
+        else:
+            assert(len(self.block_size_max) == img_dim)
+            block_size_max = self.block_size_max
+        block_num = random.randint(self.block_range[0], self.block_range[1])
+
+        for n in range(block_num):
+            block_size = [random.randint(block_size_min[i], block_size_max[i]) \
+                for i in range(img_dim)]    
+            coord_min = [random.randint(3, img_shape[1+i] - block_size[i] - 3) \
+                for i in range(img_dim)]
+            if(img_dim == 2):
+                random_block = np.random.rand(img_shape[0], block_size[0], block_size[1])
+                image[:, coord_min[0]:coord_min[0] + block_size[0], 
+                         coord_min[1]:coord_min[1] + block_size[1]] = random_block
+            else:
+                random_block = np.random.rand(img_shape[0], block_size[0], 
+                                              block_size[1], block_size[2])
+                image[:, coord_min[0]:coord_min[0] + block_size[0], 
+                         coord_min[1]:coord_min[1] + block_size[1],
+                         coord_min[2]:coord_min[2] + block_size[2]] = random_block
+        sample['image'] = image
+        return sample
+
+class OutPainting(AbstractTransform):
+    """
+    Out-painting of an input image, used for self-supervised learning 
+    """
+    def __init__(self, params):
+        super(OutPainting, self).__init__(params)
+        self.inverse  = params.get('OutPainting_inverse'.lower(), False)
+        self.prob     = params.get('OutPainting_probability'.lower(), 0.5)
+        self.block_range = params.get('OutPainting_block_range'.lower(), (1, 6))
+        self.block_size_min = params.get('OutPainting_block_size_min'.lower(), None)
+        self.block_size_max = params.get('OutPainting_block_size_max'.lower(), None)
+
+    def __call__(self, sample):
+        if(random.random() >  self.prob):
+            return sample
+
+        image= sample['image']       
+        img_shape = image.shape
+        img_dim = len(img_shape) - 1
+        assert(img_dim == 2 or img_dim == 3)
+        img_out = np.random.rand(*img_shape)
+
+        if(self.block_size_min is None):
+            block_size_min = [img_shape[1+i] - 4 * img_shape[1+i]//7 for i in range(img_dim)]
+        elif(isinstance(self.block_size_min, int)):
+            block_size_min = [self.block_size_min] * img_dim
+        else:
+            assert(len(self.block_size_min) == img_dim)
+            block_size_min = self.block_size_min
+
+        if(self.block_size_max is None):
+            block_size_max = [img_shape[1+i] - 3 * img_shape[1+i]//7 for i in range(img_dim)]
+        elif(isinstance(self.block_size_min, int)):
+            block_size_max = [self.block_size_max] * img_dim
+        else:
+            assert(len(self.block_size_max) == img_dim)
+            block_size_max = self.block_size_max
+        block_num = random.randint(self.block_range[0], self.block_range[1])
+
+        for n in range(block_num):
+            block_size = [random.randint(block_size_min[i], block_size_max[i]) \
+                for i in range(img_dim)]    
+            coord_min = [random.randint(3, img_shape[1+i] - block_size[i] - 3) \
+                for i in range(img_dim)]
+            if(img_dim == 2):
+                img_out[:, coord_min[0]:coord_min[0] + block_size[0], 
+                           coord_min[1]:coord_min[1] + block_size[1]] = \
+                    image[:, coord_min[0]:coord_min[0] + block_size[0], 
+                             coord_min[1]:coord_min[1] + block_size[1]]
+            else:
+                img_out[:, coord_min[0]:coord_min[0] + block_size[0], 
+                           coord_min[1]:coord_min[1] + block_size[1],
+                           coord_min[2]:coord_min[2] + block_size[2]] = \
+                    image[:, coord_min[0]:coord_min[0] + block_size[0], 
+                             coord_min[1]:coord_min[1] + block_size[1],
+                             coord_min[2]:coord_min[2] + block_size[2]] 
+        sample['image'] = img_out
+        return sample
+
+class InOutPainting(AbstractTransform):
+    """
+    Apply in-painting or out-patining randomly. They are mutually exclusive.
+    """
+    def __init__(self, params):
+        super(InOutPainting, self).__init__(params)
+        self.inverse  = params.get('InOutPainting_inverse'.lower(), False)
+        self.prob     = params.get('InOutPainting_probability'.lower(), 0.5)
+        self.in_prob  = params.get('InPainting_probability'.lower(), 0.5)
+        params['InPainting_probability']  = 1.0
+        params['outPainting_probability'] = 1.0
+        self.inpaint  = InPainting(params)
+        self.outpaint = OutPainting(params)
+
+    def __call__(self, sample):
+        if(random.random() >  self.prob):
+            return sample
+        if(random.random() < self.in_prob):
+            sample = self.inpaint(sample)
+        else:
+            sample = self.outpaint(sample)
         return sample
\ No newline at end of file
diff --git a/pymic/transform/rescale.py b/pymic/transform/rescale.py
index 04dd458..2a671fd 100644
--- a/pymic/transform/rescale.py
+++ b/pymic/transform/rescale.py
@@ -1,9 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division
 
-import torch
 import json
-import math
 import random
 import numpy as np
 from scipy import ndimage
@@ -85,28 +83,32 @@ class RandomRescale(AbstractTransform):
     The arguments should be written in the `params` dictionary, and it has the
     following fields:
 
-    :param `RandomRescale_lower_bound`: (list/tuple or int) 
+    :param `RandomRescale_lower_bound`: (list/tuple or float) 
         Desired minimal rescale ratio. If tuple/list, the length should be 3 or 2.
-    :param `RandomRescale_upper_bound`: (list/tuple or int) 
+    :param `RandomRescale_upper_bound`: (list/tuple or float) 
         Desired maximal rescale ratio. If tuple/list, the length should be 3 or 2.
+    :param `RandomRescale_probability`: (optional, float) 
+        The probability of applying RandomRescale. Default is 0.5.
     :param `RandomRescale_inverse`: (optional, bool) 
         Is inverse transform needed for inference. Default is `True`.
     """
     def __init__(self, params):
-        """
-        ratio0 (tuple/list or int): Desired minimal rescale ratio. 
-            If tuple/list, the length should be 3 or 2.
-        ratio1 (tuple/list or int): Desired maximal rescale ratio. 
-            If tuple/list, the length should be 3 or 2.
-        """
         super(RandomRescale, self).__init__(params)
         self.ratio0 = params["RandomRescale_lower_bound".lower()]
         self.ratio1 = params["RandomRescale_upper_bound".lower()]
+        self.prob   = params.get('RandomRescale_probability'.lower(), 0.5)
         self.inverse     = params.get("RandomRescale_inverse".lower(), True)
         assert isinstance(self.ratio0, (float, list, tuple))
         assert isinstance(self.ratio1, (float, list, tuple))
 
     def __call__(self, sample):
+        # if(random.random() > self.prob):
+        #     print("rescale not started")
+        #     sample['RandomRescale_triggered'] = False
+        #     return sample
+        # else:
+        #     print("rescale started")
+        #     sample['RandomRescale_triggered'] = True
         image = sample['image']
         input_shape = image.shape
         input_dim   = len(input_shape) - 1
@@ -117,13 +119,13 @@ def __call__(self, sample):
             scale = [self.ratio0[i] + random.random()*(self.ratio1[i] - self.ratio0[i]) \
                 for i in range(len(self.ratio0))]
         else:
-            scale = [self.ratio0 + random.random()*(self.ratio1 - self.ratio0) \
-                for i in range(input_dim)]
+            scale = self.ratio0 + random.random()*(self.ratio1 - self.ratio0)
+            scale = [scale] * input_dim
         scale = [1.0] + scale
         image_t = ndimage.interpolation.zoom(image, scale, order = 1)
 
         sample['image'] = image_t
-        sample['RandomRescale_origin_shape'] = json.dumps(input_shape)
+        sample['RandomRescale_Param'] = json.dumps(input_shape)
         if('label' in sample and self.task == 'segmentation'):
             label = sample['label']
             label = ndimage.interpolation.zoom(label, scale, order = 0)
@@ -136,11 +138,13 @@ def __call__(self, sample):
         return sample
 
     def inverse_transform_for_prediction(self, sample):
-        if(isinstance(sample['RandomRescale_origin_shape'], list) or \
-            isinstance(sample['RandomRescale_origin_shape'], tuple)):
-            origin_shape = json.loads(sample['RandomRescale_origin_shape'][0])
+        if(not sample['RandomRescale_triggered']):
+            return sample
+        if(isinstance(sample['RandomRescale_Param'], list) or \
+            isinstance(sample['RandomRescale_Param'], tuple)):
+            origin_shape = json.loads(sample['RandomRescale_Param'][0])
         else:
-            origin_shape = json.loads(sample['RandomRescale_origin_shape'])
+            origin_shape = json.loads(sample['RandomRescale_Param'])
         origin_dim   = len(origin_shape) - 1
         predict = sample['predict']
         input_shape = predict.shape
diff --git a/pymic/transform/rotate.py b/pymic/transform/rotate.py
index ba2d655..2aa06d4 100644
--- a/pymic/transform/rotate.py
+++ b/pymic/transform/rotate.py
@@ -1,9 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division
 
-import torch
 import json
-import math
 import random
 import numpy as np
 from scipy import ndimage
@@ -27,6 +25,8 @@ class RandomRotate(AbstractTransform):
     :param `RandomRotate_angle_range_w`: (list/tuple or None) 
         Rotation angle (degree) range along width axis (y-z plane), e.g., (-90, 90).
         If None, no rotation along this axis. Only used for 3D images. 
+    :param `RandomRotate_probability`: (optional, float) 
+        The probability of applying RandomRotate. Default is 0.5.
     :param `RandomRotate_inverse`: (optional, bool) 
         Is inverse transform needed for inference. Default is `True`.
     """
@@ -35,6 +35,7 @@ def __init__(self, params):
         self.angle_range_d  = params['RandomRotate_angle_range_d'.lower()]
         self.angle_range_h  = params['RandomRotate_angle_range_h'.lower()]
         self.angle_range_w  = params['RandomRotate_angle_range_w'.lower()]
+        self.prob = params.get('RandomRotate_probability'.lower(), 0.5)
         self.inverse = params.get('RandomRotate_inverse'.lower(), True)
 
     def __apply_transformation(self, image, transform_param_list, order = 1):
@@ -50,6 +51,11 @@ def __apply_transformation(self, image, transform_param_list, order = 1):
         return image
 
     def __call__(self, sample):
+        # if(random.random() > self.prob):
+        #     sample['RandomRotate_triggered'] = False
+        #     return sample
+        # else:
+        #     sample['RandomRotate_triggered'] = True
         image = sample['image']
         input_shape = image.shape
         input_dim = len(input_shape) - 1
@@ -66,7 +72,9 @@ def __call__(self, sample):
                 angle_w = np.random.uniform(self.angle_range_w[0], self.angle_range_w[1])
                 transform_param_list.append([angle_w, (-2, -3)])
         assert(len(transform_param_list) > 0)
-
+        # select a random transform from the possible list rather than 
+        # use a combination for higher efficiency
+        transform_param_list = [random.choice(transform_param_list)]
         sample['RandomRotate_Param'] = json.dumps(transform_param_list)
         image_t = self.__apply_transformation(image, transform_param_list, 1)
         sample['image'] = image_t
@@ -79,6 +87,8 @@ def __call__(self, sample):
         return sample
 
     def  inverse_transform_for_prediction(self, sample):
+        if(not sample['RandomRotate_triggered']):
+            return sample
         if(isinstance(sample['RandomRotate_Param'], list) or \
             isinstance(sample['RandomRotate_Param'], tuple)):
             transform_param_list = json.loads(sample['RandomRotate_Param'][0]) 
diff --git a/pymic/transform/trans_dict.py b/pymic/transform/trans_dict.py
index 8da3b38..bc72c93 100644
--- a/pymic/transform/trans_dict.py
+++ b/pymic/transform/trans_dict.py
@@ -30,14 +30,15 @@
 """
 from __future__ import print_function, division
 from pymic.transform.intensity import  *
-from pymic.transform.flip import RandomFlip
-from pymic.transform.pad import Pad
-from pymic.transform.rotate import RandomRotate
-from pymic.transform.rescale import Rescale, RandomRescale  
+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.label_convert import *  
 
 TransformDict = {
     'ChannelWiseThreshold': ChannelWiseThreshold,
@@ -47,9 +48,13 @@
     'GrayscaleToRGB': GrayscaleToRGB,
     'GammaCorrection': GammaCorrection,
     'GaussianNoise': GaussianNoise,
+    'InPainting': InPainting,
+    'InOutPainting': InOutPainting,
     'LabelConvert': LabelConvert,
     'LabelConvertNonzero': LabelConvertNonzero,
     'LabelToProbability': LabelToProbability,
+    'LocalShuffling': LocalShuffling,
+    'NonLinearTransform': NonLinearTransform,
     'NormalizeWithMeanStd': NormalizeWithMeanStd,
     'NormalizeWithMinMax': NormalizeWithMinMax,
     'NormalizeWithPercentiles': NormalizeWithPercentiles,
@@ -57,9 +62,11 @@
     'RandomCrop': RandomCrop,
     'RandomResizedCrop': RandomResizedCrop,
     'RandomRescale': RandomRescale,
+    'RandomTranspose': RandomTranspose,
     'RandomFlip': RandomFlip,
     'RandomRotate': RandomRotate,
     'ReduceLabelDim': ReduceLabelDim,
     'Rescale': Rescale,
+    'OutPainting': OutPainting,
     'Pad': Pad,
 }
diff --git a/pymic/transform/transpose.py b/pymic/transform/transpose.py
new file mode 100644
index 0000000..9c73bda
--- /dev/null
+++ b/pymic/transform/transpose.py
@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import json
+import random
+import numpy as np
+from pymic.transform.abstract_transform import AbstractTransform
+
+
+class RandomTranspose(AbstractTransform):
+    """
+    Random transpose for 3D volumes. Assume the input has a shape of [C, D, H, W], the
+    output shape will be of [C, D, H, W], [C, W, H, D] or [C, H, D, W]
+
+    The arguments should be written in the `params` dictionary, and it has the
+    following fields:
+
+    :param `RandomTranspose_inverse`: (optional, bool) 
+        Is inverse transform needed for inference. Default is `True`.
+    """
+    def __init__(self, params): 
+        super(RandomTranspose, self).__init__(params)
+        self.inverse = params.get('RandomTranspose_inverse'.lower(), True)
+
+    def __call__(self, sample):
+        image = sample['image']
+        input_shape = image.shape
+        input_dim = len(input_shape) - 1
+        assert(input_dim == 3)
+
+        rand_num = random.random()
+        if(rand_num < 0.4):
+            transpose_axis = None 
+        elif(rand_num < 0.7):
+            transpose_axis = [0, 3, 2, 1]
+        else:
+            transpose_axis = [0, 2, 1, 3]
+        sample['RandomTranspose_Param'] = json.dumps(transpose_axis) 
+        if(transpose_axis is not None):
+            image_t = np.transpose(image, transpose_axis)
+            sample['image'] = image_t
+            if('label' in sample and self.task == 'segmentation'):
+                sample['label'] = np.transpose(sample['label'] , transpose_axis)
+            if('pixel_weight' in sample and self.task == 'segmentation'):
+                sample['pixel_weight'] = np.transpose(sample['pixel_weight'] , transpose_axis)
+            
+        return sample
+
+    def  inverse_transform_for_prediction(self, sample):
+        if(isinstance(sample['RandomTranspose_Param'], list) or \
+            isinstance(sample['RandomTranspose_Param'], tuple)):
+            transpose_axis = json.loads(sample['RandomTranspose_Param'][0]) 
+        else:
+            transpose_axis = json.loads(sample['RandomTranspose_Param'])
+        if(transpose_axis is not None):
+            sample['predict']  = np.transpose(sample['predict'] , transpose_axis)
+        return sample
\ No newline at end of file
diff --git a/pymic/util/__init__.py b/pymic/util/__init__.py
index e69de29..72b8078 100644
--- a/pymic/util/__init__.py
+++ b/pymic/util/__init__.py
@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+from . import *
\ No newline at end of file
diff --git a/pymic/util/evaluation_cls.py b/pymic/util/evaluation_cls.py
index 29e3494..af11a17 100644
--- a/pymic/util/evaluation_cls.py
+++ b/pymic/util/evaluation_cls.py
@@ -90,7 +90,7 @@ def binary_evaluation(config):
     for i in range(len(gt_items)):
         assert(gt_items.iloc[i, 0] == prob_items.iloc[i, 0])
     
-    gt_data  = np.asarray(gt_items.iloc[:, 1])
+    gt_data   = np.asarray(gt_items.iloc[:, -1])
     prob_data = np.asarray(prob_items.iloc[:, 1:])
     score_list = []
     for metric in metric_list:
diff --git a/pymic/util/evaluation_seg.py b/pymic/util/evaluation_seg.py
index ec02297..ba04a73 100644
--- a/pymic/util/evaluation_seg.py
+++ b/pymic/util/evaluation_seg.py
@@ -6,11 +6,7 @@
 import csv
 import os
 import sys
-import math
 import pandas as pd
-import random
-import GeodisTK
-import configparser
 import numpy as np
 from scipy import ndimage
 from pymic.io.image_read_write import *
@@ -90,7 +86,7 @@ def get_edge_points(img):
         strt = ndimage.generate_binary_structure(2,1)
     else:
         strt = ndimage.generate_binary_structure(3,1)
-    ero  = ndimage.morphology.binary_erosion(img, strt)
+    ero  = ndimage.binary_erosion(img, strt)
     edge = np.asarray(img, np.uint8) - np.asarray(ero, np.uint8) 
     return edge 
 
@@ -114,14 +110,9 @@ def binary_hd95(s, g, spacing = None):
         spacing = [1.0] * image_dim
     else:
         assert(image_dim == len(spacing))
-    img = np.zeros_like(s)
-    if(image_dim == 2):
-        s_dis = GeodisTK.geodesic2d_raster_scan(img, s_edge, 0.0, 2)
-        g_dis = GeodisTK.geodesic2d_raster_scan(img, g_edge, 0.0, 2)
-    elif(image_dim ==3):
-        s_dis = GeodisTK.geodesic3d_raster_scan(img, s_edge, spacing, 0.0, 2)
-        g_dis = GeodisTK.geodesic3d_raster_scan(img, g_edge, spacing, 0.0, 2)
-
+    s_dis = ndimage.distance_transform_edt(1-s_edge, sampling = spacing)
+    g_dis = ndimage.distance_transform_edt(1-g_edge, sampling = spacing)
+  
     dist_list1 = s_dis[g_edge > 0]
     dist_list1 = sorted(dist_list1)
     dist1 = dist_list1[int(len(dist_list1)*0.95)]
@@ -150,13 +141,8 @@ def binary_assd(s, g, spacing = None):
         spacing = [1.0] * image_dim
     else:
         assert(image_dim == len(spacing))
-    img = np.zeros_like(s)
-    if(image_dim == 2):
-        s_dis = GeodisTK.geodesic2d_raster_scan(img, s_edge, 0.0, 2)
-        g_dis = GeodisTK.geodesic2d_raster_scan(img, g_edge, 0.0, 2)
-    elif(image_dim ==3):
-        s_dis = GeodisTK.geodesic3d_raster_scan(img, s_edge, spacing, 0.0, 2)
-        g_dis = GeodisTK.geodesic3d_raster_scan(img, g_edge, spacing, 0.0, 2)
+    s_dis = ndimage.distance_transform_edt(1-s_edge, sampling = spacing)
+    g_dis = ndimage.distance_transform_edt(1-g_edge, sampling = spacing)
 
     ns = s_edge.sum()
     ng = g_edge.sum()
@@ -260,14 +246,15 @@ def evaluation(config):
     Run evaluation of segmentation results based on a configuration dictionary `config`.
     The following fields should be provided in `config`:
 
-    :param metric: (str) The metric for evaluation. 
+    :param metric_list: (list) The list of metrics for evaluation. 
         The metric options are {`dice`, `iou`, `assd`, `hd95`, `rve`, `volume`}. 
     :param label_list: (list) The list of labels for evaluation. 
     :param label_fuse: (option, bool) If true, fuse the labels in the `label_list`
         as the foreground, and other labels as the background. Default is False. 
     :param organ_name: (str) The name of the organ for segmentation.
     :param ground_truth_folder_root: (str) The root dir of ground truth images. 
-    :param segmentation_folder_root: (str) The root dir of segmentation images. 
+    :param segmentation_folder_root: (str or list) The root dir of segmentation images. 
+        When a list is given, each list element should be the root dir of the results of one method. 
     :param evaluation_image_pair: (str) The csv file that provide the segmentation 
         images and the corresponding ground truth images. 
     :param ground_truth_label_convert_source: (optional, list) The list of source
@@ -280,7 +267,7 @@ def evaluation(config):
         labels for label conversion in the segmentation.     
     """
     
-    metric = config['metric']
+    metric_list = config['metric_list']
     label_list = config['label_list']
     label_fuse = config.get('label_fuse', False)
     organ_name = config['organ_name']
@@ -295,60 +282,62 @@ def evaluation(config):
     segmentation_label_convert_target = config.get('segmentation_label_convert_target', None)
 
     image_items = pd.read_csv(image_pair_csv)
-    item_num = len(image_items)
-    for seg_root_n in seg_root:
-        score_all_data = []
-        name_score_list= []
-        for i in range(item_num):
-            gt_name  = image_items.iloc[i, 0]
-            seg_name = image_items.iloc[i, 1]
-            # seg_name = seg_name.replace(".nii.gz", "_pred.nii.gz")
-            gt_full_name  = gt_root  + '/' + gt_name
-            seg_full_name = seg_root_n + '/' + seg_name
-            
-            s_dict = load_image_as_nd_array(seg_full_name)
-            g_dict = load_image_as_nd_array(gt_full_name)
-            s_volume = s_dict["data_array"]; s_spacing = s_dict["spacing"]
-            g_volume = g_dict["data_array"]; g_spacing = g_dict["spacing"]
-            # for dim in range(len(s_spacing)):
-            #     assert(s_spacing[dim] == g_spacing[dim])
-            if((ground_truth_label_convert_source is not None) and \
-                ground_truth_label_convert_target is not None):
-                g_volume = convert_label(g_volume, ground_truth_label_convert_source, \
-                    ground_truth_label_convert_target)
-
-            if((segmentation_label_convert_source is not None) and \
-                segmentation_label_convert_target is not None):
-                s_volume = convert_label(s_volume, segmentation_label_convert_source, \
-                    segmentation_label_convert_target)
-
-            score_vector = get_multi_class_evaluation_score(s_volume, g_volume, label_list, 
-                label_fuse, s_spacing, metric )
-            if(len(label_list) > 1):
-                score_vector.append(np.asarray(score_vector).mean())
-            score_all_data.append(score_vector)
-            name_score_list.append([seg_name] + score_vector)
-            print(seg_name, score_vector)
-        score_all_data = np.asarray(score_all_data)
-        score_mean = score_all_data.mean(axis = 0)
-        score_std  = score_all_data.std(axis = 0)
-        name_score_list.append(['mean'] + list(score_mean))
-        name_score_list.append(['std'] + list(score_std))
+    item_num    = len(image_items)
     
-        # save the result as csv 
-        score_csv = "{0:}/{1:}_{2:}_all.csv".format(seg_root_n, organ_name, metric)
-        with open(score_csv, mode='w') as csv_file:
-            csv_writer = csv.writer(csv_file, delimiter=',', 
-                            quotechar='"',quoting=csv.QUOTE_MINIMAL)
-            head = ['image'] + ["class_{0:}".format(i) for i in label_list]
-            if(len(label_list) > 1):
-                head = head + ["average"]
-            csv_writer.writerow(head)
-            for item in name_score_list:
-                csv_writer.writerow(item)
-
-        print("{0:} mean ".format(metric), score_mean)
-        print("{0:} std  ".format(metric), score_std) 
+    for seg_root_n in seg_root: # for each segmentation method
+        for metric in metric_list:
+            score_all_data = []
+            name_score_list= []
+            for i in range(item_num):
+                gt_name  = image_items.iloc[i, 0]
+                seg_name = image_items.iloc[i, 1]
+                # seg_name = seg_name.replace(".nii.gz", "_pred.nii.gz")
+                gt_full_name  = gt_root  + '/' + gt_name
+                seg_full_name = seg_root_n + '/' + seg_name
+                
+                s_dict = load_image_as_nd_array(seg_full_name)
+                g_dict = load_image_as_nd_array(gt_full_name)
+                s_volume = s_dict["data_array"]; s_spacing = s_dict["spacing"]
+                g_volume = g_dict["data_array"]; g_spacing = g_dict["spacing"]
+                # for dim in range(len(s_spacing)):
+                #     assert(s_spacing[dim] == g_spacing[dim])
+                if((ground_truth_label_convert_source is not None) and \
+                    ground_truth_label_convert_target is not None):
+                    g_volume = convert_label(g_volume, ground_truth_label_convert_source, \
+                        ground_truth_label_convert_target)
+
+                if((segmentation_label_convert_source is not None) and \
+                    segmentation_label_convert_target is not None):
+                    s_volume = convert_label(s_volume, segmentation_label_convert_source, \
+                        segmentation_label_convert_target)
+
+                score_vector = get_multi_class_evaluation_score(s_volume, g_volume, label_list, 
+                    label_fuse, s_spacing, metric )
+                if(len(label_list) > 1):
+                    score_vector.append(np.asarray(score_vector).mean())
+                score_all_data.append(score_vector)
+                name_score_list.append([seg_name] + score_vector)
+                print(seg_name, score_vector)
+            score_all_data = np.asarray(score_all_data)
+            score_mean = score_all_data.mean(axis = 0)
+            score_std  = score_all_data.std(axis = 0)
+            name_score_list.append(['mean'] + list(score_mean))
+            name_score_list.append(['std'] + list(score_std))
+        
+            # save the result as csv 
+            score_csv = "{0:}/{1:}_{2:}_all.csv".format(seg_root_n, organ_name, metric)
+            with open(score_csv, mode='w') as csv_file:
+                csv_writer = csv.writer(csv_file, delimiter=',', 
+                                quotechar='"',quoting=csv.QUOTE_MINIMAL)
+                head = ['image'] + ["class_{0:}".format(i) for i in label_list]
+                if(len(label_list) > 1):
+                    head = head + ["average"]
+                csv_writer.writerow(head)
+                for item in name_score_list:
+                    csv_writer.writerow(item)
+
+            print("{0:} mean ".format(metric), score_mean)
+            print("{0:} std  ".format(metric), score_std) 
 
 def main():
     """
diff --git a/pymic/util/general.py b/pymic/util/general.py
index 75b6af1..075d2e1 100644
--- a/pymic/util/general.py
+++ b/pymic/util/general.py
@@ -9,6 +9,25 @@ def keyword_match(a,b):
     """
     return a.lower() == b.lower()
 
+def tensor_shape_match(a,b):
+    """
+    Test if two tensors have the same shape"""
+    shape_a = list(a.shape)
+    shape_b = list(b.shape)
+    len_a  = len(shape_a)
+    len_b  = len(shape_b)
+    if(len_a != len_b):
+        return False 
+    elif(len_a == 0):
+        return True 
+    else:
+        for i in range(len_a):
+            if(shape_a[i] != shape_b[i]):
+                return False 
+        return True 
+
+
+
 def get_one_hot_seg(label, class_num):
     """
     Convert a segmentation label to one-hot.
@@ -29,4 +48,52 @@ def get_one_hot_seg(label, class_num):
     one_hot = one_hot.view(*size)
     one_hot = torch.transpose(one_hot, 1, -1)
     one_hot = torch.squeeze(one_hot, -1)
-    return one_hot
\ No newline at end of file
+    return one_hot
+
+def mixup(inputs, labels):
+    """Shuffle a minibatch and do linear interpolation between images and labels.
+    Both classification and segmentation labels are supported. The targets should
+    be one-hot labels.
+    
+    :param inputs: a tensor of input images with size N X C0 x H x W.
+    :param labels: a tensor of one-hot labels. The shape is N X C for classification
+        tasks, and N X C X H X W for segmentation tasks. 
+    """
+    input_shape = list(inputs.shape)
+    label_shape = list(labels.shape)
+    img_dim     = len(input_shape) - 2
+    N = input_shape[0] # batch size
+    C = label_shape[1] # class number
+    rp1 = torch.randperm(N)
+    inputs1 = inputs[rp1]
+    labels1 = labels[rp1]
+    
+    rp2 = torch.randperm(N)
+    inputs2 = inputs[rp2]
+    labels2 = labels[rp2]
+
+    a = np.random.beta(1, 1, [N, 1])
+    if(img_dim == 2):
+        b = np.tile(a[..., None, None], [1] + input_shape[1:])
+    elif(img_dim == 3):
+        b = np.tile(a[..., None, None, None], [1] + input_shape[1:])
+    else:
+        raise ValueError("MixUp only supports 2D and 3D images, but the " +
+            "input image has {0:} dimensions".format(img_dim))
+
+    inputs1 = inputs1 * torch.from_numpy(b).float()
+    inputs2 = inputs2 * torch.from_numpy(1 - b).float()
+    inputs_mix = inputs1 + inputs2
+
+    if(len(label_shape) == 2): # for classification tasks
+        c = np.tile(a, [1, C])
+    elif(img_dim == 2):        # for 2D segmentation tasks
+        c = np.tile(a[..., None, None], [1] + label_shape[1:])
+    else:                      # for 3D segmentation tasks
+        c = np.tile(a[..., None, None, None], [1] + label_shape[1:])
+    
+    labels1 = labels1 * torch.from_numpy(c).float()
+    labels2 = labels2 * torch.from_numpy(1 - c).float()
+    labels_mix = labels1 + labels2
+
+    return inputs_mix, labels_mix
diff --git a/pymic/util/parse_config.py b/pymic/util/parse_config.py
index 5e81312..232762f 100644
--- a/pymic/util/parse_config.py
+++ b/pymic/util/parse_config.py
@@ -112,9 +112,12 @@ def synchronize_config(config):
 
 def logging_config(config):
     for section in config:
-        for key in config[section]:
-            value = config[section][key]
-            logging.info("{0:} {1:} = {2:}".format(section, key, value))
+        if(isinstance(config[section], dict)):
+            for key in config[section]:
+                value = config[section][key]
+                logging.info("{0:} {1:} = {2:}".format(section, key, value))
+        else:
+            logging.info("{0:} = {1:}".format(section, config[section]))
 
 if __name__ == "__main__":
     print(is_int('555'))
diff --git a/pymic/util/post_process.py b/pymic/util/post_process.py
index a0a9dff..a889b23 100644
--- a/pymic/util/post_process.py
+++ b/pymic/util/post_process.py
@@ -43,6 +43,7 @@ def __call__(self, seg):
                 seg_c  = np.asarray(seg == c, np.uint8)
                 seg_c  = get_largest_k_components(seg_c) 
                 output = output + seg_c * c
+            seg = output
         return seg
 
 PostProcessDict = {
diff --git a/pyproject.toml b/pyproject.toml
deleted file mode 100644
index 4bbff29..0000000
--- a/pyproject.toml
+++ /dev/null
@@ -1,8 +0,0 @@
-[build-system]
-requires = ["flit_core >=3.2,<4"]
-build-backend = "flit_core.buildapi"
-
-[project]
-name = "PyMIC"
-authors = [{name = "Graziella", email = "graziella@lumache"}]
-dynamic = ["version", "description"]
diff --git a/requirements.txt b/requirements.txt
index c8cd562..cac47f3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+h5py
 matplotlib>=3.1.2
 numpy>=1.17.4
 pandas>=0.25.3
@@ -5,7 +6,7 @@ scikit-image>=0.16.2
 scikit-learn>=0.22
 scipy>=1.3.3
 SimpleITK>=2.0.0
-tensorboard>=2.1.0
-tensorboardX>=1.9
-torch>=1.7.1
-torchvision>=0.8.2
+tensorboard
+tensorboardX
+torch>=1.1.12
+torchvision>=0.13.0
diff --git a/setup.py b/setup.py
index ce7271b..36daf9a 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name    = 'PYMIC',
-    version = "0.3.0",
+    version = "0.4.0",
     author  ='PyMIC Consortium',
     author_email = 'wguotai@gmail.com',
     description  = description,
@@ -20,6 +20,18 @@
     url      = 'https://github.com/HiLab-git/PyMIC',
     license  = 'Apache 2.0',
     packages = setuptools.find_packages(),
+    install_requires=[
+        "h5py",
+        "matplotlib>=3.1.2",
+        "numpy>=1.17.4",
+        "pandas>=0.25.3",
+        "scikit-image>=0.16.2",
+        "scikit-learn>=0.22",
+        "scipy>=1.3.3",
+        "SimpleITK>=2.0.0",
+        "tensorboard",
+        "tensorboardX",
+    ],
     classifiers=[
         'License :: OSI Approved :: Apache Software License',
         'Programming Language :: Python',
@@ -29,10 +41,8 @@
     python_requires = '>=3.6',
     entry_points = {
         'console_scripts': [
-            'pymic_run  = pymic.net_run.net_run:main',
-            'pymic_ssl  = pymic.net_run_ssl.ssl_main:main',
-            'pymic_wsl  = pymic.net_run_wsl.wsl_main:main',
-            'pymic_nll  = pymic.net_run_nll.nll_main:main',
+            'pymic_train = pymic.net_run.train:main',
+            'pymic_test  = pymic.net_run.predict:main',
             'pymic_eval_cls = pymic.util.evaluation_cls:main',
             'pymic_eval_seg = pymic.util.evaluation_seg:main'
         ],


Test patch:
diff --git a/pymic/test/test_assd.py b/pymic/test/test_assd.py
new file mode 100644
index 0000000..35c1804
--- /dev/null
+++ b/pymic/test/test_assd.py
@@ -0,0 +1,37 @@
+from scipy import ndimage
+from PIL import Image
+import numpy as np 
+import SimpleITK as sitk
+import matplotlib.pyplot as plt
+from pymic.util.evaluation_seg import get_edge_points
+
+def test_assd_2d():
+    img_name = "/home/x/projects/PyMIC_project/PyMIC_examples/PyMIC_data/JSRT/label/JPCLN001.png"
+    img = Image.open(img_name)
+    img_array = np.asarray(img)
+    img_edge = get_edge_points(img_array > 0)
+    s_dis = ndimage.distance_transform_edt(1-img_edge)
+    plt.subplot(1,2,1)
+    plt.imshow(img_edge)
+    plt.subplot(1,2,2)
+    plt.imshow(s_dis)
+    plt.show()
+
+def test_assd_3d():
+    img_name = "/home/x/projects/PyMIC_project/PyMIC_examples/seg_ssl/ACDC/result/unet2d_baseline/patient001_frame01.nii.gz"
+    img_obj = sitk.ReadImage(img_name) 
+    spacing = img_obj.GetSpacing()
+    spacing = spacing[::-1]
+    img_data = sitk.GetArrayFromImage(img_obj)
+    print(img_data.shape)
+    print(spacing)
+    img_edge = get_edge_points(img_data > 0)
+    s_dis = ndimage.distance_transform_edt(1-img_edge, sampling=spacing)
+    dis_obj = sitk.GetImageFromArray(s_dis)
+    dis_obj.CopyInformation(img_obj)
+    sitk.WriteImage(dis_obj, "test_dis.nii.gz")
+
+
+
+if __name__ == "__main__":
+    test_assd_3d()
\ No newline at end of file
diff --git a/pymic/test/test_nifty_dataset.py b/pymic/test/test_nifty_dataset.py
index 581bc52..b5c4660 100644
--- a/pymic/test/test_nifty_dataset.py
+++ b/pymic/test/test_nifty_dataset.py
@@ -1,50 +1,68 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division
+import sys
+import numpy as np 
+from pymic.util.parse_config import *
+from pymic.net_run.agent_cls import ClassificationAgent
+from pymic.net_run.agent_seg import SegmentationAgent
+import SimpleITK as sitk
 
-import os
-import torch
-import pandas as pd
-import numpy as np
-from skimage import io, transform
-from torch.utils.data import Dataset, DataLoader
-from torchvision import transforms, utils
-from pymic.io.image_read_write import *
-from pymic.io.nifty_dataset import NiftyDataset
-from pymic.io.transform3d import *
+def save_array_as_nifty_volume(data, image_name, reference_name = None):
+    """
+    Save a numpy array as nifty image
 
-if __name__ == "__main__":
-    root_dir = '/home/guotai/data/brats/BraTS2018_Training'
-    csv_file = '/home/guotai/projects/torch_brats/brats/config/brats18_train_train.csv'
-    
-    crop1 = CropWithBoundingBox(start = None, output_size = [4, 144, 180, 144])
-    norm  = ChannelWiseNormalize(mean = None, std = None, zero_to_random = True)
-    labconv = LabelConvert([0, 1, 2, 4], [0, 1, 2, 3])
-    crop2 = RandomCrop([128, 128, 128])
-    rescale =Rescale([64, 64, 64])
-    transform_list = [crop1, norm, labconv, crop2,rescale, ToTensor()]
-    transformed_dataset = NiftyDataset(root_dir=root_dir,
-                                    csv_file  = csv_file,
-                                    modal_num = 4,
-                                    transform = transforms.Compose(transform_list))
-    dataloader = DataLoader(transformed_dataset, batch_size=4,
-                            shuffle=True, num_workers=4)
-    # Helper function to show a batch
+    :param data:  (numpy.ndarray) A numpy array with shape [Depth, Height, Width].
+    :param image_name:  (str) The ouput file name.
+    :param reference_name:  (str) File name of the reference image of which 
+        meta information is used.
+    """
+    img = sitk.GetImageFromArray(data)
+    if(reference_name is not None):
+        img_ref = sitk.ReadImage(reference_name)
+        #img.CopyInformation(img_ref)
+        img.SetSpacing(img_ref.GetSpacing())
+        img.SetOrigin(img_ref.GetOrigin())
+        img.SetDirection(img_ref.GetDirection())
+    sitk.WriteImage(img, image_name)
 
+def main():
+    """
+    The main function for running a network for training or inference.
+    """
+    if(len(sys.argv) < 3):
+        print('Number of arguments should be 3. e.g.')
+        print('python test_nifty_dataset.py train config.cfg')
+        exit()
+    stage    = str(sys.argv[1])
+    cfg_file = str(sys.argv[2])
+    config   = parse_config(cfg_file)
+    config   = synchronize_config(config)
+    # task     = config['dataset']['task_type']
+    # assert task in ['cls', 'cls_nexcl', 'seg']
+    # if(task == 'cls' or task == 'cls_nexcl'):
+    #     agent = ClassificationAgent(config, stage)
+    # else:
+    #     agent = SegmentationAgent(config, stage)
+    agent = SegmentationAgent(config, stage)
+    agent.create_dataset()
+    data_loader = agent.train_loader if stage == "train" else agent.test_loader
+    it = 0
+    for data in data_loader:
+        inputs      = agent.convert_tensor_type(data['image'])
+        labels_prob = agent.convert_tensor_type(data['label_prob']) 
+        for i in range(inputs.shape[0]):
+            image_i = inputs[i][0]
+            label_i = np.argmax(labels_prob[i], axis = 0)
+            print(image_i.shape, label_i.shape)
+            image_name = "temp/image_{0:}_{1:}.nii.gz".format(it, i)
+            label_name = "temp/label_{0:}_{1:}.nii.gz".format(it, i)
+            save_array_as_nifty_volume(image_i, image_name, reference_name = None)
+            save_array_as_nifty_volume(label_i, label_name, reference_name = None)
+        it = it + 1
+        if(it == 10):
+            break
 
-    for i_batch, sample_batched in enumerate(dataloader):
-        print(i_batch, sample_batched['image'].size(),
-            sample_batched['label'].size())
+if __name__ == "__main__":
+    main()
+    
 
-        # # observe 4th batch and stop.
-        modals = ['flair', 't1ce', 't1', 't2']
-        if i_batch == 0:
-            image =  sample_batched['image'].numpy()
-            label =  sample_batched['label'].numpy()
-            for i in range(image.shape[0]):
-                for mod in range(4):
-                    image_i = image[i][mod]
-                    label_i = label[i][0]
-                    image_name = "temp/image_{0:}_{1:}.nii.gz".format(i, modals[mod])
-                    label_name = "temp/label_{0:}.nii.gz".format(i)
-                    save_array_as_nifty_volume(image_i, image_name, reference_name = None)
-                    save_array_as_nifty_volume(label_i, label_name, reference_name = None)
\ No newline at end of file