Basic train utils

Author: qubvel

segmentation_models_pytorch/utils/functions.py

@@ -1,11 +1,11 @@
 import torch
 
 
-def iou(outputs, targets, eps=1e-7, threshold=None, activation='sigmoid'):
+def iou(pr, gt, eps=1e-7, threshold=None, activation='sigmoid'):
     """
     Args:
-        outputs (torch.Tensor): A list of predicted elements
-        targets (torch.Tensor):  A list of elements that are to be predicted
+        pr (torch.Tensor): A list of predicted elements
+        gt (torch.Tensor):  A list of elements that are to be predicted
         eps (float): epsilon to avoid zero division
         threshold: threshold for outputs binarization
     Returns:
@@ -20,16 +20,54 @@ def iou(outputs, targets, eps=1e-7, threshold=None, activation='sigmoid'):
         activation_fn = torch.nn.Softmax2d()
     else:
         raise NotImplementedError(
-            "Dice is only implemented for sigmoid and softmax2d"
+            "Activation implemented for sigmoid and softmax2d"
         )
 
-    outputs = activation_fn(outputs)
+    pr = activation_fn(pr)
 
     if threshold is not None:
-        outputs = (outputs > threshold).float()
+        pr = (pr > threshold).float()
 
-    intersection = torch.sum(targets * outputs)
-    union = torch.sum(targets) + torch.sum(outputs) - intersection + eps
+    intersection = torch.sum(gt * pr)
+    union = torch.sum(gt) + torch.sum(pr) - intersection + eps
     return (intersection + eps) / union
 
-jaccard = iou
+jaccard = iou
+
+
+def f_score(pr, gt, beta=1, eps=1e-7, threshold=None, activation='sigmoid'):
+    """
+    Args:
+        pr (torch.Tensor): A list of predicted elements
+        gt (torch.Tensor):  A list of elements that are to be predicted
+        eps (float): epsilon to avoid zero division
+        threshold: threshold for outputs binarization
+    Returns:
+        float: IoU (Jaccard) score
+    """
+
+    if activation is None or activation == "none":
+        activation_fn = lambda x: x
+    elif activation == "sigmoid":
+        activation_fn = torch.nn.Sigmoid()
+    elif activation == "softmax2d":
+        activation_fn = torch.nn.Softmax2d()
+    else:
+        raise NotImplementedError(
+            "Activation implemented for sigmoid and softmax2d"
+        )
+
+    pr = activation_fn(pr)
+
+    if threshold is not None:
+        pr = (pr > threshold).float()
+
+
+    tp = torch.sum(gt * pr)
+    fp = torch.sum(pr) - tp
+    fn = torch.sum(gt) - tp
+
+    score = ((1 + beta ** 2) * tp + eps) \
+            / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + eps)
+
+    return score

segmentation_models_pytorch/utils/losses.py

@@ -3,7 +3,6 @@
 
 
 class JaccardLoss(nn.Module):
-
     __name__ = 'jaccard_loss'
 
     def __init__(self, eps=1e-7, activation='sigmoid'):
@@ -12,11 +11,22 @@ def __init__(self, eps=1e-7, activation='sigmoid'):
         self.eps = eps
 
     def forward(self, y_pr, y_gt):
-        return F.iou(y_pr, y_gt, self.eps, threshold=None, activation=self.activation)
+        return 1 - F.jaccard(y_pr, y_gt, eps=self.eps, threshold=None, activation=self.activation)
 
 
-class BCEJaccardLoss(JaccardLoss):
+class DiceLoss(nn.Module):
+    __name__ = 'dice_loss'
+
+    def __init__(self, eps=1e-7, activation='sigmoid'):
+        super().__init__()
+        self.activation = activation
+        self.eps = eps
+
+    def forward(self, y_pr, y_gt):
+        return 1 - F.f_score(y_pr, y_gt, beta=1., eps=self.eps, threshold=None, activation=self.activation)
+
 
+class BCEJaccardLoss(JaccardLoss):
     __name__ = 'bce_jaccard_loss'
 
     def __init__(self, eps=1e-7, activation='sigmoid'):
@@ -26,4 +36,17 @@ def __init__(self, eps=1e-7, activation='sigmoid'):
     def forward(self, y_pr, y_gt):
         jaccard = super().forward(y_pr, y_gt)
         bce = self.bce(y_pr, y_gt)
-        return jaccard + bce
+        return jaccard + bce
+
+
+class BCEDiceLoss(DiceLoss):
+    __name__ = 'bce_dice_loss'
+
+    def __init__(self, eps=1e-7, activation='sigmoid'):
+        super().__init__(eps, activation)
+        self.bce = nn.BCEWithLogitsLoss(reduction='mean')
+
+    def forward(self, y_pr, y_gt):
+        dice = super().forward(y_pr, y_gt)
+        bce = self.bce(y_pr, y_gt)
+        return dice + bce

segmentation_models_pytorch/utils/metrics.py

@@ -14,3 +14,18 @@ def __init__(self, eps=1e-7, threshold=0.5, activation='sigmoid'):
 
     def forward(self, y_pr, y_gt):
         return F.iou(y_pr, y_gt, self.eps, self.threshold, self.activation)
+
+
+class FscoreMetric(nn.Module):
+
+    __name__ = 'f-score'
+
+    def __init__(self, beta=1, eps=1e-7, threshold=0.5, activation='sigmoid'):
+        super().__init__()
+        self.activation = activation
+        self.eps = eps
+        self.threshold = threshold
+        self.beta = beta
+
+    def forward(self, y_pr, y_gt):
+        return F.f_score(y_pr, y_gt, self.beta, self.eps, self.threshold, self.activation)

segmentation_models_pytorch/utils/train.py

@@ -14,8 +14,16 @@ def __init__(self, model, loss, metrics, stage_name, device='cpu', verbose=True)
         self.verbose = verbose
         self.device = device
 
+        self._to_device()
+
+    def _to_device(self):
+        self.model.to(self.device)
+        self.loss.to(self.device)
+        for metric in self.metrics:
+            metric.to(self.device)
+
     def _format_logs(self, logs):
-        str_logs = ['{}_{} - {:.4}'.format(self.stage_name, k, v) for k, v in logs.items()]
+        str_logs = ['{} - {:.4}'.format(k, v) for k, v in logs.items()]
         s = ', '.join(str_logs)
         return s
 
@@ -25,33 +33,30 @@ def batch_update(self, x, y):
     def on_epoch_start(self):
         pass
 
-    def run(self, dataloder):
+    def run(self, dataloader):
 
         self.on_epoch_start()
 
         logs = {}
         loss_meter = AverageValueMeter()
-        metrics_meters = {m.__name__: AverageValueMeter() for m in self.metrics}
+        metrics_meters = {metric.__name__: AverageValueMeter() for metric in self.metrics}
 
-        with tqdm(dataloder, desc=self.stage_name, file=sys.stdout, disable=not (self.verbose)) as iterator:
+        with tqdm(dataloader, desc=self.stage_name, file=sys.stdout, disable=not (self.verbose)) as iterator:
             for x, y in iterator:
                 x, y = x.to(self.device), y.to(self.device)
                 loss, y_pred = self.batch_update(x, y)
 
-                # update loss meter
+                # update loss logs
                 loss_value = loss.cpu().detach().numpy()
                 loss_meter.add(loss_value)
+                loss_logs = {self.loss.__name__: loss_meter.mean}
+                logs.update(loss_logs)
 
-                # update metrics meters
+                # update metrics logs
                 for metric_fn in self.metrics:
                     metric_value = metric_fn(y_pred, y).cpu().detach().numpy()
                     metrics_meters[metric_fn.__name__].add(metric_value)
-
-                # create_logs
-                loss_logs = {'loss': loss_meter.mean}
                 metrics_logs = {k: v.mean for k, v in metrics_meters.items()}
-
-                logs.update(loss_logs)
                 logs.update(metrics_logs)
 
                 if self.verbose:
@@ -105,4 +110,4 @@ def batch_update(self, x, y):
         with torch.no_grad():
             prediction = self.model.forward(x)
             loss = self.loss(prediction, y)
-            return loss, prediction
+        return loss, prediction