support mixup

Author: taigw

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):

pymic/net_run/agent_cls.py

@@ -9,6 +9,7 @@
 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
@@ -17,6 +18,7 @@
 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.*')
 
@@ -111,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)
@@ -129,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
@@ -140,8 +144,11 @@ 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
@@ -174,7 +181,7 @@ 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()
                 # forward + backward + optimize

pymic/net_run/agent_seg.py

@@ -12,6 +12,7 @@
 import torch.optim as optim
 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
@@ -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
 
 class SegmentationAgent(NetRunAgent):
     def __init__(self, config, stage = 'train'):
@@ -120,6 +122,7 @@ def set_postprocessor(self, postprocessor):
     def training(self):
         class_num   = self.config['network']['class_num']
         iter_valid  = self.config['training']['iter_valid']
+        mixup_prob  = self.config['training'].get('mixup_probability', 0.5)
         train_loss  = 0
         train_dice_list = []
         self.net.train()
@@ -132,7 +135,9 @@ def training(self):
             # get the inputs
             inputs      = self.convert_tensor_type(data['image'])
             labels_prob = self.convert_tensor_type(data['label_prob'])                 
-            
+            if(random() < mixup_prob):
+                inputs, labels_prob = mixup(inputs, labels_prob) 
+                   
             # # for debug
             # for i in range(inputs.shape[0]):
             #     image_i = inputs[i][0]

pymic/util/general.py

@@ -29,4 +29,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
+    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