update loss function

Author: taigw

examples/JSRT/config/evaluation.cfg

@@ -11,3 +11,7 @@ ground_truth_label_convert_source = None
 ground_truth_label_convert_target = None
 segmentation_label_convert_source = None
 segmentation_label_convert_target = None
+
+
+
+

examples/JSRT/config/train_test.cfg

@@ -51,11 +51,12 @@ bilinear      = True
 
 [training]
 # device name" cuda:n or cpu
-device_name = cuda:0
+device_name = cuda:1
 
 batch_size    = 4
-loss_function = dice_loss
-class_weight  = [0.5, 1.0]
+loss_type     = DiceLoss
+DiceLoss_enable_pixel_weight = False
+DiceLoss_enable_class_weight = False
 
 # for optimizers
 optimizer     = Adam

examples/fetal_hc/config/train_test.cfg

@@ -62,8 +62,9 @@ bilinear      = True
 device_name = cuda:0
 
 batch_size    = 4
-loss_function = dice_loss
-class_weight  = [0.5, 1.0]
+loss_type     = DiceLoss
+DiceLoss_enable_pixel_weight = False
+DiceLoss_enable_class_weight = False
 
 # for optimizers
 optimizer     = Adam

examples/prostate/config/train_test.cfg

@@ -59,8 +59,9 @@ trilinear     = True
 device_name = cuda:0
 
 batch_size  = 4
-loss_function = dice_loss
-class_weight  = [0.5, 1.0]
+loss_type     = DiceLoss
+DiceLoss_enable_pixel_weight = False
+DiceLoss_enable_class_weight = False
 
 # for optimizers
 optimizer     = Adam

pymic/loss/__init__.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+from pymic.loss.util import reshape_tensor_to_2D
+
+class CrossEntropyLoss(nn.Module):
+    def __init__(self, params):
+        super(CrossEntropyLoss, self).__init__()
+        self.enable_pix_weight = params['CrossEntropyLoss_Enable_Pixel_Weight'.lower()]
+        self.enable_cls_weight = params['CrossEntropyLoss_Enable_Class_Weight'.lower()]
+    
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        pix_w   = loss_input_dict['pixel_weight']
+        cls_w   = loss_input_dict['class_weight']
+        softmax = loss_input_dict['softmax']
+
+        if(softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+        predict = reshape_tensor_to_2D(predict)
+        soft_y  = reshape_tensor_to_2D(soft_y)
+
+        ce = - soft_y* torch.log(predict)
+        if(self.enable_cls_weight):
+            if(cls_w is None):
+                raise ValueError("Class weight is enabled but not defined")
+            ce = torch.sum(ce * cls_w, dim = 1)
+        else:
+            ce = torch.sum(ce, dim = 1)
+        if(self.enable_pix_weight):
+            if(pix_w is None):
+                raise ValueError("Pixel weight is enabled but not defined")
+            pix_w = reshape_tensor_to_2D(pix_w)
+            ce    = torch.sum(ce * pix_w) / torch.sum(pix_w)
+        else:
+            ce = torch.mean(ce)  
+        return ce

pymic/loss/ce.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+from pymic.loss.util import reshape_tensor_to_2D, get_classwise_dice
+
+class DiceLoss(nn.Module):
+    def __init__(self, params):
+        super(DiceLoss, self).__init__()
+        self.enable_pix_weight = params['DiceLoss_Enable_Pixel_Weight'.lower()]
+        self.enable_cls_weight = params['DiceLoss_Enable_Class_Weight'.lower()]
+
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        pix_w   = loss_input_dict['pixel_weight']
+        cls_w   = loss_input_dict['class_weight']
+        softmax = loss_input_dict['softmax']
+
+        if(softmax):
+            predict = nn.Softmax(dim = 1)(predict)
+        predict = reshape_tensor_to_2D(predict)
+        soft_y  = reshape_tensor_to_2D(soft_y) 
+
+        if(self.enable_pix_weight):
+            if(pix_w is None):
+                raise ValueError("Pixel weight is enabled but not defined")
+            pix_w = reshape_tensor_to_2D(pix_w)
+        dice_score = get_classwise_dice(predict, soft_y, pix_w)
+        if(self.enable_cls_weight):
+            if(cls_w is None):
+                raise ValueError("Class weight is enabled but not defined")
+            weighted_dice = dice_score * cls_w
+            avg_dice =  weighted_dice.sum() / cls_w.sum()
+        else:
+            avg_dice = torch.mean(dice_score)   
+        dice_loss  = 1.0 - avg_dice
+        return dice_loss
+
+class MultiScaleDiceLoss(nn.Module):
+    def __init__(self, params):
+        super(MultiScaleDiceLoss, self).__init__()
+        self.enable_pix_weight  = params['MultiScaleDiceLoss_Enable_Pixel_Weight'.lower()]
+        self.enable_cls_weight  = params['MultiScaleDiceLoss_Enable_Class_Weight'.lower()]
+        self.multi_scale_weight = params['MultiScaleDiceLoss_Scale_Weight'.lower()]
+        dice_params = {'DiceLoss_Enable_Pixel_Weight'.lower(): self.enable_pix_weight,
+            'DiceLoss_Enable_Class_Weight'.lower(): self.enable_cls_weight}
+        self.base_loss = DiceLoss(dice_params)
+
+    def forward(self, loss_input_dict):
+        predict = loss_input_dict['prediction']
+        soft_y  = loss_input_dict['ground_truth']
+        pix_w   = loss_input_dict['pixel_weight']
+        cls_w   = loss_input_dict['class_weight']
+        softmax = loss_input_dict['softmax']
+        if(isinstance(predict, tuple) or isinstance(predict, list)):
+            predict_num = len(predict)
+            assert(predict_num == len(self.multi_scale_weight))
+            loss   = 0.0
+            weight = 0.0
+            interp_mode = 'trilinear' if(len(predict[0].shape) == 5) else 'bilinear'
+            for i in range(predict_num):
+                soft_y_temp = nn.functional.interpolate(soft_y, 
+                    size = list(predict[i].shape)[2:], mode = interp_mode)
+                if(pix_w is not None):
+                    pix_w_temp  = nn.functional.interpolate(pix_w, 
+                        size = list(predict[i].shape)[2:], mode = interp_mode)
+                else:
+                    pix_w_temp = None
+                temp_loss_dict = {}
+                temp_loss_dict['prediction'] = predict[i]
+                temp_loss_dict['ground_truth'] = soft_y_temp
+                temp_loss_dict['pixel_weight'] = pix_w_temp
+                temp_loss_dict['class_weight'] = cls_w
+                temp_loss_dict['softmax'] = softmax
+                temp_loss = self.base_loss(temp_loss_dict)
+                loss      = loss + temp_loss * self.multi_scale_weight[i]
+                weight    = weight + self.multi_scale_weight[i]
+            loss = loss/weight
+        else:
+            loss = self.base_loss(loss_input_dict)
+        return loss

pymic/loss/dice.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+from pymic.loss.ce import CrossEntropyLoss
+from pymic.loss.dice import DiceLoss, MultiScaleDiceLoss
+
+loss_dict = {'CrossEntropyLoss': CrossEntropyLoss,
+    'DiceLoss': DiceLoss,
+    'MultiScaleDiceLoss': MultiScaleDiceLoss}
+
+def get_loss(params):
+    loss_type = params['loss_type']
+    if(loss_type in loss_dict):
+        loss_obj = loss_dict[loss_type](params)
+    else:
+        raise ValueError("Undefined loss type {0:}".format(loss_type))
+    return loss_obj

pymic/loss/loss_factory.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+def get_soft_label(input_tensor, num_class, data_type = 'float'):
+    """
+        convert a label tensor to soft label 
+        input_tensor: tensor with shae [B, 1, D, H, W]
+        output_tensor: shape [B, num_class, D, H, W]
+    """
+    tensor_list = []
+    for i in range(num_class):
+        temp_prob = input_tensor == i*torch.ones_like(input_tensor)
+        tensor_list.append(temp_prob)
+    output_tensor = torch.cat(tensor_list, dim = 1)
+    if(data_type == 'float'):
+        output_tensor = output_tensor.float()
+    elif(data_type == 'double'):
+        output_tensor = output_tensor.double()
+    else:
+        raise ValueError("data type can only be float and double: {0:}".format(data_type))
+
+    return output_tensor
+
+def reshape_tensor_to_2D(x):
+    """
+    reshape input variables of shape [B, C, D, H, W] to [voxel_n, C]
+    """
+    tensor_dim = len(x.size())
+    num_class  = list(x.size())[1]
+    if(tensor_dim == 5):
+        x_perm  = x.permute(0, 2, 3, 4, 1)
+    elif(tensor_dim == 4):
+        x_perm  = x.permute(0, 2, 3, 1)
+    else:
+        raise ValueError("{0:}D tensor not supported".format(tensor_dim))
+    
+    y = torch.reshape(x_perm, (-1, num_class)) 
+    return y 
+
+def reshape_prediction_and_ground_truth(predict, soft_y):
+    """
+    reshape input variables of shape [B, C, D, H, W] to [voxel_n, C]
+    """
+    tensor_dim = len(predict.size())
+    num_class  = list(predict.size())[1]
+    if(tensor_dim == 5):
+        soft_y  = soft_y.permute(0, 2, 3, 4, 1)
+        predict = predict.permute(0, 2, 3, 4, 1)
+    elif(tensor_dim == 4):
+        soft_y  = soft_y.permute(0, 2, 3, 1)
+        predict = predict.permute(0, 2, 3, 1)
+    else:
+        raise ValueError("{0:}D tensor not supported".format(tensor_dim))
+    
+    predict = torch.reshape(predict, (-1, num_class)) 
+    soft_y  = torch.reshape(soft_y,  (-1, num_class))
+      
+    return predict, soft_y
+
+def get_classwise_dice(predict, soft_y, pix_w = None):
+    """
+    get dice scores for each class in predict (after softmax) and soft_y
+    """
+    
+    if(pix_w is None):
+        y_vol = torch.sum(soft_y,  dim = 0)
+        p_vol = torch.sum(predict, dim = 0)
+        intersect = torch.sum(soft_y * predict, dim = 0)
+    else:
+        y_vol = torch.sum(soft_y * pix_w,  dim = 0)
+        p_vol = torch.sum(predict * pix_w, dim = 0)
+        intersect = torch.sum(soft_y * predict * pix_w, dim = 0)
+    dice_score = (2.0 * intersect + 1e-5)/ (y_vol + p_vol + 1e-5)
+    return dice_score