New losses: focal loss & generalised dice loss

README.md

@@ -47,7 +47,7 @@ Please find below a description of each parameter:
 - `batch_norm_momentum`: float (e.g. 0.1).
 - `num_epochs`: int.
 - `initial_lr`: initial learning rate.
-- `loss`: choice between 'dice' and 'cross_entropy'. Note: Please use 'cross_entropy' when comparing `Unet` vs. `MixedUp-Unet`.
+- `loss`: dictionary with a key `'name'` for the choice between `'dice'`, `'focal'`, `'focal_dice'`, `'gdl'` and `'cross_entropy'` and a (optional) key `'params'` (e.g.`{"name": "focal", "params": {"gamma": 0.5}}`.
 - `log_directory`: folder name where log files are saved.
 - `film_layers`: indicates on which layer(s) of the U-net you want to apply a FiLM modulation: list of 8 elements (because Unet has 8 layers), set to 0 for no FiLM modulation, set 1 otherwise. Note: When running `Unet` or `MixedUp-Unet`, please fill this list with zeros only.
 - `mixup_bool`: indicates if mixup is applied to the training data (choice: `false` or `true`). Note: Please use `false` when comparing `Unet` vs. `FiLMed-Unet`.

config/config.json

@@ -13,9 +13,9 @@
         "batch_norm_momentum": 0.1,
         "num_epochs": 100,
         "initial_lr": 0.001,
-	"loss": "dice",
+	"loss": {"name": "dice"},
         "log_directory": "log_sc",
-        "film_layers": [1, 0, 0, 0, 0, 0, 0, 0],
+        "film_layers": [0, 0, 0, 0, 0, 0, 0, 0],
 	"mixup_bool": false,
 	"mixup_alpha": 2,
 	"metadata": "contrast",

config/config_gm.json

@@ -13,7 +13,7 @@
 	"batch_norm_momentum": 0.1,
 	"num_epochs": 500,
 	"initial_lr": 0.001,
-        "loss": "dice",
+        "loss": {"name": "dice"},
 	"log_directory": "log_gm",
         "film_layers": [0, 0, 0, 0, 0, 0, 0, 0],
 	"mixup_bool": false,

config/config_sctTesting.json

@@ -8,16 +8,16 @@
         "contrast_balance": {},
         "contrast_test": ["T2w", "T1w", "acq-axmid00100_T2w", "acq-axbottom_T2w", "acq-sagtsp_T2w", "acq-sagstir_T2w", "acq-axtop_T2w", "acq-ax00014_T2w", "acq-axtsp_T2w", "acq-axtop00100_T2w", "acq-axcsp_T2w", "acq-sup_T2star", "acq-ax_T2w", "acq-sagcsp_T2w", "acq-axmid00005_T2w", "acq-sagstirtsp_T2w", "acq-ax00012_T2w", "acq-axlow_T2w", "T2star", "acq-sagstircsp_T2w", "acq-sag_T2w", "acq-sag00015_T2w", "acq-inf_T2star"],
         "center_test": [],
-        "batch_size": 18,
+        "batch_size": 32,
         "dropout_rate": 0.3,
         "batch_norm_momentum": 0.1,
         "num_epochs": 100,
-        "initial_lr": 0.001,
-        "loss": "dice",
+        "initial_lr": 0.0001,
+        "loss": {"name": "focal_dice", "params": {"gamma": 0.4, "alpha": 10}},
         "log_directory": "log_large",
-        "film_layers": [1, 0, 0, 0, 0, 0, 0, 0],
+        "film_layers": [0, 0, 0, 0, 0, 0, 0, 0],
 	"mixup_bool": false,
 	"mixup_alpha": 2,
-	"metadata": "contrast",
+	"metadata": "without",
         "debugging": false
 }

config/config_small.json

@@ -2,7 +2,7 @@
 	"command": "train",
 	"gpu": 4,
         "gt_suffix": "_seg-manual",
-	"bids_path_train": "../duke/projects/ivado-medical-imaging/spineGeneric_201907041011/result/",
+	"bids_path": "../duke/projects/ivado-medical-imaging/spineGeneric_201907041011/result/",
 	"random_seed": 1313,
 	"contrast_train_validation": ["T1w", "T2w", "T2star", "acq-MToff_MTS", "acq-MTon_MTS", "acq-T1w_MTS"],
 	"contrast_balance": {},
@@ -13,11 +13,11 @@
 	"batch_norm_momentum": 0.1,
 	"num_epochs": 100,
 	"initial_lr": 0.001,
-        "loss": "dice",
+        "loss": {"name": "dice"},
 	"log_directory": "log_sc_small",
-	"film_layers": [1, 0, 0, 0, 0, 0, 0, 0],
-	"mixup_bool": false,
+	"film_layers": [0, 0, 0, 0, 0, 0, 0, 0],
+	"mixup_bool": true,
 	"mixup_alpha": 2,
-	"metadata": "contrast",
+	"metadata": "without",
         "debugging": false
 }

ivadomed/losses.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def dice_loss(input, target):
+    # input = torch.sigmoid(input)
+    smooth = 1.0
+
+    iflat = input.view(-1)
+    tflat = target.view(-1)
+    intersection = (iflat * tflat).sum()
+
+    return ((2.0 * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth))
+
+
+class FocalLoss(nn.Module):
+    """
+    Focal Loss: https://arxiv.org/abs/1708.02002
+    """
+
+    def __init__(self, gamma):
+        super().__init__()
+        self.gamma = gamma
+
+    def forward(self, input, target):
+        if not (target.size() == input.size()):
+            raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
+
+        max_val = (-input).clamp(min=0)
+        loss = input - input * target + max_val + ((-max_val).exp() + (-input - max_val).exp()).log()
+
+        # This gives us the log sigmoid of 1-p if y is 0 and of p if y is 1
+        invprobs = F.logsigmoid(-input * (target * 2 - 1))
+        loss = (invprobs * self.gamma).exp() * loss
+
+        # Note: works in log space to be numerically stable (ie to avoid NaNs when training).
+        return loss.mean()
+
+
+class FocalDiceLoss(nn.Module):
+    """
+    Motivated by https://arxiv.org/pdf/1809.00076.pdf
+    :param alpha: to bring the dice and focal losses at similar scale.
+    :param gamma: gamma value used in the focal loss.
+    """
+    def __init__(self, alpha, gamma):
+        super().__init__()
+        self.alpha = alpha
+        self.focal = FocalLoss(gamma)
+
+    def forward(self, input, target):
+        loss = self.alpha * self.focal(input, target) - torch.log(dice_loss(input, target))
+        return loss.mean()
+
+
+class GeneralizedDiceLoss(nn.Module):
+    """
+    Generalized Dice Loss: https://arxiv.org/pdf/1707.03237
+    """
+    def __init__(self, epsilon=1e-5):
+        super(GeneralizedDiceLoss, self).__init__()
+        self.epsilon = epsilon
+
+    def forward(self, input, target):
+        if not (target.size() == input.size()):
+            raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
+
+        input = input.view(-1)
+        target = target.view(-1)
+
+        target = target.float()
+        target_sum = target.sum(-1)
+        class_weights = nn.Parameter(1. / (target_sum * target_sum).clamp(min=self.epsilon))
+
+        intersect = (input * target).sum(-1) * class_weights
+        intersect = intersect.sum()
+
+        denominator = ((input + target).sum(-1) * class_weights).sum()
+
+        return 1. - 2. * intersect / denominator.clamp(min=self.epsilon)
+

ivadomed/main.py

@@ -5,6 +5,8 @@
 import shutil
 import random
 import joblib
+from math import exp
+import numpy as np
 
 import torch
 import torch.nn as nn
@@ -16,7 +18,6 @@
 
 from medicaltorch import transforms as mt_transforms
 from medicaltorch import datasets as mt_datasets
-from medicaltorch import losses as mt_losses
 from medicaltorch import filters as mt_filters
 from medicaltorch import metrics as mt_metrics
 
@@ -26,10 +27,9 @@
 
 from ivadomed import loader as loader
 from ivadomed import models
+from ivadomed import losses
 from ivadomed.utils import *
 
-import numpy as np
-
 cudnn.benchmark = True
 
 
@@ -176,24 +176,36 @@ def cmd_train(context):
     var_contrast_list = []
 
     # Loss
-    if context["loss"] in ["dice", "cross_entropy"]:
-        if context["loss"] == "cross_entropy":
+    if context["loss"]["name"] in ["dice", "cross_entropy", "focal", "gdl", "focal_dice"]:
+        if context["loss"]["name"] == "cross_entropy":
             loss_fct = nn.BCELoss()
+        elif context["loss"]["name"] == "focal":
+            loss_fct = losses.FocalLoss(gamma=context["loss"]["params"]["gamma"])
+            print("\nLoss function: {}, with gamma={}.\n".format(context["loss"]["name"], context["loss"]["params"]["gamma"]))
+        elif context["loss"]["name"] == "gdl":
+            loss_fct = losses.GeneralizedDiceLoss()
+        elif context["loss"]["name"] == "focal_dice":
+            loss_fct = losses.FocalDiceLoss(gamma=context["loss"]["params"]["gamma"], alpha=context["loss"]["params"]["alpha"])
+            print("\nLoss function: {}, with gamma={} and alpha={}.\n".format(context["loss"]["name"], context["loss"]["params"]["gamma"], context["loss"]["params"]["alpha"]))
+            focal_loss_fct = losses.FocalLoss(gamma=context["loss"]["params"]["gamma"]) # for tuning alpha
+
+        if not context["loss"]["name"].startswith("focal"):
+            print("\nLoss function: {}.\n".format(context["loss"]["name"]))
+
     else:
-        print("Unknown Loss function, please choose between 'dice' or 'cross_entropy'")
+        print("Unknown Loss function, please choose between 'dice', 'focal', 'focal_dice', 'gdl' or 'cross_entropy'")
         exit()
 
     # Training loop -----------------------------------------------------------
     best_validation_loss = float("inf")
-    bce_loss = nn.BCELoss()
     for epoch in tqdm(range(1, num_epochs+1), desc="Training"):
         start_time = time.time()
 
         lr = scheduler.get_lr()[0]
         writer.add_scalar('learning_rate', lr, epoch)
 
         model.train()
-        train_loss_total = 0.0
+        train_loss_total, dice_train_loss_total, focal_train_loss_total = 0.0, 0.0, 0.0
         num_steps = 0
         for i, batch in enumerate(train_loader):
             input_samples, gt_samples = batch["input"], batch["gt"]
@@ -207,7 +219,6 @@ def cmd_train(context):
                     mixup_folder = os.path.join(context["log_directory"], 'mixup')
                     if not os.path.isdir(mixup_folder):
                         os.makedirs(mixup_folder)
-                    print(lambda_tensor.data.numpy()[0])
                     random_idx = np.random.randint(0, input_samples.size()[0])
                     val_gt = np.unique(gt_samples.data.numpy()[random_idx,0,:,:])
                     mixup_fname_pref = os.path.join(mixup_folder, str(i).zfill(3)+'_'+str(lambda_tensor.data.numpy()[0])+'_'+str(random_idx).zfill(3)+'.png')
@@ -234,10 +245,15 @@ def cmd_train(context):
             else:
                 preds = model(var_input)
 
-            if context["loss"] == "dice":
-                loss = mt_losses.dice_loss(preds, var_gt)
+            if context["loss"]["name"] == "dice":
+                loss = - losses.dice_loss(preds, var_gt)
             else:
                 loss = loss_fct(preds, var_gt)
+                if context["loss"]["name"] == "focal_dice":
+                    focal_train_loss_total += focal_loss_fct(preds, var_gt).item()
+                    dice_train_loss_total += torch.log(losses.dice_loss(preds, var_gt)).item()
+                else:
+                    dice_train_loss_total += losses.dice_loss(preds, var_gt).item()
             train_loss_total += loss.item()
 
             optimizer.zero_grad()
@@ -267,13 +283,23 @@ def cmd_train(context):
         train_loss_total_avg = train_loss_total / num_steps
 
         tqdm.write(f"Epoch {epoch} training loss: {train_loss_total_avg:.4f}.")
+        if context["loss"]["name"] == 'focal_dice':
+            focal_train_loss_total_avg = focal_train_loss_total / num_steps
+            log_dice_train_loss_total_avg = dice_train_loss_total / num_steps
+            dice_train_loss_total_avg = exp(log_dice_train_loss_total_avg)
+            tqdm.write(f"\tFocal training loss: {focal_train_loss_total_avg:.4f}.")
+            tqdm.write(f"\tLog Dice training loss: {log_dice_train_loss_total_avg:.4f}.")
+            tqdm.write(f"\tDice training loss: {dice_train_loss_total_avg:.4f}.")
+        elif context["loss"]["name"] != 'dice':
+            dice_train_loss_total_avg = dice_train_loss_total / num_steps
+            tqdm.write(f"\tDice training loss: {dice_train_loss_total_avg:.4f}.")
 
         # Validation loop -----------------------------------------------------
         model.eval()
-        val_loss_total = 0.0
+        val_loss_total, dice_val_loss_total, focal_val_loss_total = 0.0, 0.0, 0.0
         num_steps = 0
 
-        metric_fns = [mt_metrics.dice_score,
+        metric_fns = [dice_score,  # from ivadomed/utils.py
                       mt_metrics.hausdorff_score,
                       mt_metrics.precision_score,
                       mt_metrics.recall_score,
@@ -304,11 +330,15 @@ def cmd_train(context):
                 else:
                     preds = model(var_input)
 
-                # loss = mt_losses.dice_loss(preds, var_gt)
-                if context["loss"] == "dice":
-                    loss = mt_losses.dice_loss(preds, var_gt)
+                if context["loss"]["name"] == "dice":
+                    loss = - losses.dice_loss(preds, var_gt)
                 else:
                     loss = loss_fct(preds, var_gt)
+                    if context["loss"]["name"] == "focal_dice":
+                        focal_val_loss_total += focal_loss_fct(preds, var_gt).item()
+                        dice_val_loss_total += torch.log(losses.dice_loss(preds, var_gt)).item()
+                    else:
+                        dice_val_loss_total += losses.dice_loss(preds, var_gt).item()
                 val_loss_total += loss.item()
 
             # Metrics computation
@@ -369,6 +399,16 @@ def cmd_train(context):
         }, epoch)
 
         tqdm.write(f"Epoch {epoch} validation loss: {val_loss_total_avg:.4f}.")
+        if context["loss"]["name"] == 'focal_dice':
+            focal_val_loss_total_avg = focal_val_loss_total / num_steps
+            log_dice_val_loss_total_avg = dice_val_loss_total / num_steps
+            dice_val_loss_total_avg = exp(log_dice_val_loss_total_avg)
+            tqdm.write(f"\tFocal validation loss: {focal_val_loss_total_avg:.4f}.")
+            tqdm.write(f"\tLog Dice validation loss: {log_dice_val_loss_total_avg:.4f}.")
+            tqdm.write(f"\tDice validation loss: {dice_val_loss_total_avg:.4f}.")
+        elif context["loss"]["name"] != 'dice':
+            dice_val_loss_total_avg = dice_val_loss_total / num_steps
+            tqdm.write(f"\tDice validation loss: {dice_val_loss_total_avg:.4f}.")
 
         end_time = time.time()
         total_time = end_time - start_time
@@ -465,8 +505,8 @@ def cmd_test(context):
         model.cuda()
     model.eval()
 
-    metric_fns = [mt_metrics.dice_score,
-                  # mt_metrics.hausdorff_score,
+    metric_fns = [dice_score,  # from ivadomed/utils.py
+                  mt_metrics.hausdorff_score,
                   mt_metrics.precision_score,
                   mt_metrics.recall_score,
                   mt_metrics.specificity_score,

ivadomed/utils.py

@@ -7,6 +7,23 @@
 from medicaltorch import filters as mt_filters
 from medicaltorch import transforms as mt_transforms
 
+
+def dice_score(im1, im2, eps=1.0):
+    """
+    Computes the Dice coefficient between im1 and im2.
+    """
+    im1 = np.asarray(im1).astype(np.bool)
+    im2 = np.asarray(im2).astype(np.bool)
+
+    if im1.shape != im2.shape:
+        raise ValueError("Shape mismatch: im1 and im2 must have the same shape.")
+
+    im_sum = im1.sum() + im2.sum() + eps
+
+    intersection = np.logical_and(im1, im2)
+    return 2. * intersection.sum() / im_sum
+
+
 def mixup(data, targets, alpha):
     """Compute the mixup data.
     Return mixed inputs and targets, lambda.