Merge pull request #6 from qubvel/feature/train

Train utils

Author: qubvel

requirements.txt

@@ -1,2 +1,3 @@
 torchvision==0.2.2
-pretrainedmodels==0.7.4
+pretrainedmodels==0.7.4
+torchnet==0.0.4

segmentation_models_pytorch/__init__.py

@@ -3,4 +3,5 @@
 from .fpn import FPN
 from .pspnet import PSPNet
 
-from . import encoders
+from . import encoders
+from . import utils

segmentation_models_pytorch/base/encoder_decoder.py

@@ -18,16 +18,7 @@ def __init__(self, encoder, decoder, activation):
             self.activation = nn.Sigmoid()
         else:
             raise ValueError('Activation should be "sigmoid" or "softmax"')
-        
-        if encoder.pretrained:
-            self.set_preprocessing_params(
-                input_size=encoder.input_size,
-                input_space=encoder.input_space,
-                input_range=encoder.input_range,
-                mean=encoder.mean,
-                std=encoder.std,
-            )
-
+            
     def forward(self, x):
         x = self.encoder(x)
         x = self.decoder(x)

segmentation_models_pytorch/encoders/__init__.py

@@ -45,7 +45,7 @@ def get_preprocessing_fn(encoder_name, pretrained='imagenet'):
     mean = settings[pretrained].get('mean')
     std = settings[pretrained].get('std')
 
-    def _preprocess_input(x):
-        return preprocess_input(x, mean=mean, std=std, input_space=input_space, input_range=input_range)
+    def _preprocess_input(x, **kwargs):
+        return preprocess_input(x, mean=mean, std=std, input_space=input_space, input_range=input_range, **kwargs)
 
     return _preprocess_input

segmentation_models_pytorch/encoders/_preprocessing.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 
-def preprocess_input(x, mean=None, std=None, input_space='RGB', input_range=None):
+def preprocess_input(x, mean=None, std=None, input_space='RGB', input_range=None, **kwargs):
 
     if input_space == 'BGR':
         x = x[..., ::-1].copy()

segmentation_models_pytorch/utils/__init__.py

@@ -0,0 +1,3 @@
+from . import train
+from . import losses
+from . import metrics

segmentation_models_pytorch/utils/functions.py

@@ -0,0 +1,75 @@
+import torch
+
+
+def iou(pr, gt, eps=1e-7, threshold=None, activation='sigmoid'):
+    """
+    Source:
+        https://github.com/catalyst-team/catalyst/
+    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()
+
+    intersection = torch.sum(gt * pr)
+    union = torch.sum(gt) + torch.sum(pr) - intersection + eps
+    return (intersection + eps) / union
+
+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

@@ -0,0 +1,52 @@
+import torch.nn as nn
+from . import functions as F
+
+
+class JaccardLoss(nn.Module):
+    __name__ = 'jaccard_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.jaccard(y_pr, y_gt, eps=self.eps, threshold=None, activation=self.activation)
+
+
+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'):
+        super().__init__(eps, activation)
+        self.bce = nn.BCEWithLogitsLoss(reduction='mean')
+
+    def forward(self, y_pr, y_gt):
+        jaccard = super().forward(y_pr, y_gt)
+        bce = self.bce(y_pr, y_gt)
+        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

@@ -0,0 +1,31 @@
+import torch.nn as nn
+from . import functions as F
+
+
+class IoUMetric(nn.Module):
+
+    __name__ = 'iou'
+
+    def __init__(self, eps=1e-7, threshold=0.5, activation='sigmoid'):
+        super().__init__()
+        self.activation = activation
+        self.eps = eps
+        self.threshold = threshold
+
+    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

@@ -0,0 +1,113 @@
+import sys
+import torch
+from tqdm import tqdm as tqdm
+from torchnet.meter import AverageValueMeter
+
+
+class Epoch:
+
+    def __init__(self, model, loss, metrics, stage_name, device='cpu', verbose=True):
+        self.model = model
+        self.loss = loss
+        self.metrics = metrics
+        self.stage_name = stage_name
+        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(k, v) for k, v in logs.items()]
+        s = ', '.join(str_logs)
+        return s
+
+    def batch_update(self, x, y):
+        raise NotImplementedError
+
+    def on_epoch_start(self):
+        pass
+
+    def run(self, dataloader):
+
+        self.on_epoch_start()
+
+        logs = {}
+        loss_meter = AverageValueMeter()
+        metrics_meters = {metric.__name__: AverageValueMeter() for metric in self.metrics}
+
+        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 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 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)
+                metrics_logs = {k: v.mean for k, v in metrics_meters.items()}
+                logs.update(metrics_logs)
+
+                if self.verbose:
+                    s = self._format_logs(logs)
+                    iterator.set_postfix_str(s)
+
+        return logs
+
+
+class TrainEpoch(Epoch):
+
+    def __init__(self, model, loss, metrics, optimizer, device='cpu', verbose=True):
+        super().__init__(
+            model=model,
+            loss=loss,
+            metrics=metrics,
+            stage_name='train',
+            device=device,
+            verbose=verbose,
+        )
+        self.optimizer = optimizer
+
+    def on_epoch_start(self):
+        self.model.train()
+
+    def batch_update(self, x, y):
+        self.optimizer.zero_grad()
+        prediction = self.model.forward(x)
+        loss = self.loss(prediction, y)
+        loss.backward()
+        self.optimizer.step()
+        return loss, prediction
+
+
+class ValidEpoch(Epoch):
+
+    def __init__(self, model, loss, metrics, device='cpu', verbose=True):
+        super().__init__(
+            model=model,
+            loss=loss,
+            metrics=metrics,
+            stage_name='valid',
+            device=device,
+            verbose=verbose,
+        )
+
+    def on_epoch_start(self):
+        self.model.eval()
+
+    def batch_update(self, x, y):
+        with torch.no_grad():
+            prediction = self.model.forward(x)
+            loss = self.loss(prediction, y)
+        return loss, prediction