Initial train example

Author: qubvel

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/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,35 @@
+import torch
+
+
+def iou(outputs, targets, 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
+        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(
+            "Dice is only implemented for sigmoid and softmax2d"
+        )
+
+    outputs = activation_fn(outputs)
+
+    if threshold is not None:
+        outputs = (outputs > threshold).float()
+
+    intersection = torch.sum(targets * outputs)
+    union = torch.sum(targets) + torch.sum(outputs) - intersection + eps
+    return (intersection + eps) / union
+
+jaccard = iou

segmentation_models_pytorch/utils/losses.py

@@ -0,0 +1,29 @@
+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 F.iou(y_pr, y_gt, 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

segmentation_models_pytorch/utils/metrics.py

@@ -0,0 +1,16 @@
+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)

segmentation_models_pytorch/utils/train.py

@@ -0,0 +1,108 @@
+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
+
+    def _format_logs(self, logs):
+        str_logs = ['{}_{} - {:.4}'.format(self.stage_name, 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, dataloder):
+
+        self.on_epoch_start()
+
+        logs = {}
+        loss_meter = AverageValueMeter()
+        metrics_meters = {m.__name__: AverageValueMeter() for m in self.metrics}
+
+        with tqdm(dataloder, 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
+                loss_value = loss.cpu().detach().numpy()
+                loss_meter.add(loss_value)
+
+                # update metrics meters
+                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:
+                    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