Merge pull request #29 from BloodAxe/develop

PyTorch Toolbelt 0.2.1

Author: BloodAxe

README.md

@@ -125,4 +125,19 @@ merged_mask = tiler.crop_to_orignal_size(merged_mask)
 ## Advanced examples
 
 1. [Inria Sattelite Segmentation](https://github.com/BloodAxe/Catalyst-Inria-Segmentation-Example)
-1. [CamVid Semantic Segmentation](https://github.com/BloodAxe/Catalyst-CamVid-Segmentation-Example)
+1. [CamVid Semantic Segmentation](https://github.com/BloodAxe/Catalyst-CamVid-Segmentation-Example)
+
+
+## Citation
+
+```
+@misc{Khvedchenya_Eugene_2019_PyTorch_Toolbelt,
+  author = {Khvedchenya, Eugene},
+  title = {PyTorch Toolbelt},
+  year = {2019},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/BloodAxe/pytorch-toolbelt}},
+  commit = {cc5e9973cdb0dcbf1c6b6e1401bf44b9c69e13f3}
+}
+```

pytorch_toolbelt/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import absolute_import
 
-__version__ = "0.2.0"
+__version__ = "0.2.1"

pytorch_toolbelt/losses/focal.py

@@ -2,7 +2,7 @@
 
 from torch.nn.modules.loss import _Loss
 
-from .functional import sigmoid_focal_loss, reduced_focal_loss
+from .functional import focal_loss_with_logits
 
 __all__ = ["BinaryFocalLoss", "FocalLoss"]
 
@@ -31,14 +31,15 @@ def __init__(
         self.ignore_index = ignore_index
         if reduced:
             self.focal_loss = partial(
-                reduced_focal_loss,
+                focal_loss_with_logits,
+                alpha=None,
                 gamma=gamma,
                 threshold=threshold,
                 reduction=reduction,
             )
         else:
             self.focal_loss = partial(
-                sigmoid_focal_loss, gamma=gamma, alpha=alpha, reduction=reduction
+                focal_loss_with_logits, alpha=alpha, gamma=gamma, reduction=reduction
             )
 
     def forward(self, label_input, label_target):
@@ -87,7 +88,7 @@ def forward(self, label_input, label_target):
                 cls_label_target = cls_label_target[not_ignored]
                 cls_label_input = cls_label_input[not_ignored]
 
-            loss += sigmoid_focal_loss(
+            loss += focal_loss_with_logits(
                 cls_label_input, cls_label_target, gamma=self.gamma, alpha=self.alpha
             )
         return loss

pytorch_toolbelt/losses/functional.py

@@ -1,14 +1,27 @@
 import math
+from typing import Optional
 
 import torch
 import torch.nn.functional as F
 
-__all__ = ["sigmoid_focal_loss", "soft_jaccard_score", "soft_dice_score", "wing_loss"]
+__all__ = [
+    "focal_loss_with_logits",
+    "sigmoid_focal_loss",
+    "soft_jaccard_score",
+    "soft_dice_score",
+    "wing_loss",
+]
 
 
-def sigmoid_focal_loss(
-    input: torch.Tensor, target: torch.Tensor, gamma=2.0, alpha=0.25, reduction="mean"
-):
+def focal_loss_with_logits(
+    input: torch.Tensor,
+    target: torch.Tensor,
+    gamma=2.0,
+    alpha: Optional[float] = 0.25,
+    reduction="mean",
+    normalized=False,
+    threshold: Optional[float] = None,
+) -> torch.Tensor:
     """Compute binary focal loss between target and output logits.
 
     See :class:`~pytorch_toolbelt.losses.FocalLoss` for details.
@@ -23,7 +36,8 @@ def sigmoid_focal_loss(
             and :attr:`reduce` are in the process of being deprecated, and in the meantime,
             specifying either of those two args will override :attr:`reduction`.
             'batchwise_mean' computes mean loss per sample in batch. Default: 'mean'
-
+        normalized (bool): Compute normalized focal loss (https://arxiv.org/pdf/1909.07829.pdf).
+        threshold (float, optional): Compute reduced focal loss (https://arxiv.org/abs/1903.01347).
     References::
 
         https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/loss/losses.py
@@ -34,11 +48,21 @@ def sigmoid_focal_loss(
     pt = torch.exp(logpt)
 
     # compute the loss
-    loss = -((1 - pt).pow(gamma)) * logpt
+    if threshold is None:
+        focal_term = (1 - pt).pow(gamma)
+    else:
+        focal_term = ((1.0 - pt) / threshold).pow(gamma)
+        focal_term[pt < threshold] = 1
+
+    loss = -focal_term * logpt
 
     if alpha is not None:
         loss = loss * (alpha * target + (1 - alpha) * (1 - target))
 
+    if normalized:
+        norm_factor = focal_term.sum()
+        loss = loss / norm_factor
+
     if reduction == "mean":
         loss = loss.mean()
     if reduction == "sum":
@@ -49,51 +73,21 @@ def sigmoid_focal_loss(
     return loss
 
 
+# TODO: Mark as deprecated and emit warning
+sigmoid_focal_loss = focal_loss_with_logits
+
+
+# TODO: Mark as deprecated and emit warning
 def reduced_focal_loss(
     input: torch.Tensor,
     target: torch.Tensor,
     threshold=0.5,
     gamma=2.0,
     reduction="mean",
 ):
-    """Compute reduced focal loss between target and output logits.
-
-    See :class:`~pytorch_toolbelt.losses.FocalLoss` for details.
-
-    Args:
-        input: Tensor of arbitrary shape
-        target: Tensor of the same shape as input
-        reduction (string, optional): Specifies the reduction to apply to the output:
-            'none' | 'mean' | 'sum' | 'batchwise_mean'. 'none': no reduction will be applied,
-            'mean': the sum of the output will be divided by the number of
-            elements in the output, 'sum': the output will be summed. Note: :attr:`size_average`
-            and :attr:`reduce` are in the process of being deprecated, and in the meantime,
-            specifying either of those two args will override :attr:`reduction`.
-            'batchwise_mean' computes mean loss per sample in batch. Default: 'mean'
-
-    References::
-
-        https://arxiv.org/abs/1903.01347
-    """
-    target = target.type(input.type())
-
-    logpt = -F.binary_cross_entropy_with_logits(input, target, reduction="none")
-    pt = torch.exp(logpt)
-
-    # compute the loss
-    focal_reduction = ((1.0 - pt) / threshold).pow(gamma)
-    focal_reduction[pt < threshold] = 1
-
-    loss = -focal_reduction * logpt
-
-    if reduction == "mean":
-        loss = loss.mean()
-    if reduction == "sum":
-        loss = loss.sum()
-    if reduction == "batchwise_mean":
-        loss = loss.sum(0)
-
-    return loss
+    return focal_loss_with_logits(
+        input, target, alpha=None, gamma=gamma, reduction=reduction, threshold=threshold
+    )
 
 
 def soft_jaccard_score(

pytorch_toolbelt/modules/encoders.py

@@ -603,24 +603,34 @@ def __init__(
         strides: List[int],
         channels: List[int],
         layers: List[int],
+        first_avg_pool=False,
     ):
+        if layers is None:
+            layers = [1, 2, 3, 4]
+
         super().__init__(channels, strides, layers)
 
+        def except_pool(block: nn.Module):
+            del block.pool
+            return block
+
         self.layer0 = nn.Sequential(
             densenet.features.conv0, densenet.features.norm0, densenet.features.relu0
         )
-        self.pool0 = densenet.features.pool0
+
+        self.avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
+        self.pool0 = self.avg_pool if first_avg_pool else densenet.features.pool0
 
         self.layer1 = nn.Sequential(
-            densenet.features.denseblock1, densenet.features.transition1
+            densenet.features.denseblock1, except_pool(densenet.features.transition1)
         )
 
         self.layer2 = nn.Sequential(
-            densenet.features.denseblock2, densenet.features.transition2
+            densenet.features.denseblock2, except_pool(densenet.features.transition2)
         )
 
         self.layer3 = nn.Sequential(
-            densenet.features.denseblock3, densenet.features.transition3
+            densenet.features.denseblock3, except_pool(densenet.features.transition3)
         )
 
         self.layer4 = nn.Sequential(densenet.features.denseblock4)
@@ -650,54 +660,60 @@ def forward(self, x):
             if layer == self.layer0:
                 # Fist maxpool operator is not a part of layer0 because we want that layer0 output to have stride of 2
                 output = self.pool0(output)
+            else:
+                output = self.avg_pool(output)
+
             input = output
 
         # Return only features that were requested
         return _take(output_features, self._layers)
 
 
 class DenseNet121Encoder(DenseNetEncoder):
-    def __init__(self, layers=None, pretrained=True, memory_efficient=False):
-        if layers is None:
-            layers = [1, 2, 3, 4]
+    def __init__(
+        self, layers=None, pretrained=True, memory_efficient=False, first_avg_pool=False
+    ):
         densenet = densenet121(pretrained=pretrained, memory_efficient=memory_efficient)
         strides = [2, 4, 8, 16, 32]
         channels = [64, 128, 256, 512, 1024]
-        super().__init__(densenet, strides, channels, layers)
+        super().__init__(densenet, strides, channels, layers, first_avg_pool)
 
 
 class DenseNet161Encoder(DenseNetEncoder):
-    def __init__(self, layers=None, pretrained=True, memory_efficient=False):
-        if layers is None:
-            layers = [1, 2, 3, 4]
+    def __init__(
+        self, layers=None, pretrained=True, memory_efficient=False, first_avg_pool=False
+    ):
         densenet = densenet161(pretrained=pretrained, memory_efficient=memory_efficient)
         strides = [2, 4, 8, 16, 32]
         channels = [96, 192, 384, 1056, 2208]
-        super().__init__(densenet, strides, channels, layers)
+        super().__init__(densenet, strides, channels, layers, first_avg_pool)
 
 
 class DenseNet169Encoder(DenseNetEncoder):
-    def __init__(self, layers=None, pretrained=True, memory_efficient=False):
-        if layers is None:
-            layers = [1, 2, 3, 4]
+    def __init__(
+        self, layers=None, pretrained=True, memory_efficient=False, first_avg_pool=False
+    ):
         densenet = densenet169(pretrained=pretrained, memory_efficient=memory_efficient)
         strides = [2, 4, 8, 16, 32]
         channels = [64, 128, 256, 640, 1664]
-        super().__init__(densenet, strides, channels, layers)
+        super().__init__(densenet, strides, channels, layers, first_avg_pool)
 
 
 class DenseNet201Encoder(DenseNetEncoder):
-    def __init__(self, layers=None, pretrained=True, memory_efficient=False):
-        if layers is None:
-            layers = [1, 2, 3, 4]
+    def __init__(
+        self, layers=None, pretrained=True, memory_efficient=False, first_avg_pool=False
+    ):
         densenet = densenet201(pretrained=pretrained, memory_efficient=memory_efficient)
         strides = [2, 4, 8, 16, 32]
         channels = [64, 128, 256, 896, 1920]
-        super().__init__(densenet, strides, channels, layers)
+        super().__init__(densenet, strides, channels, layers, first_avg_pool)
 
 
 class EfficientNetEncoder(EncoderModule):
     def __init__(self, efficientnet, filters, strides, layers):
+        if layers is None:
+            layers = [1, 2, 4, 6]
+
         super().__init__(filters, strides, layers)
 
         self.stem = efficientnet.stem
@@ -736,7 +752,7 @@ def forward(self, x):
 
 
 class EfficientNetB0Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b0(num_classes=1, **kwargs),
             [16, 24, 40, 80, 112, 192, 320],
@@ -746,7 +762,7 @@ def __init__(self, layers=[1, 2, 4, 6], **kwargs):
 
 
 class EfficientNetB1Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b1(num_classes=1, **kwargs),
             [16, 24, 40, 80, 112, 192, 320],
@@ -756,7 +772,7 @@ def __init__(self, layers=[1, 2, 4, 6], **kwargs):
 
 
 class EfficientNetB2Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b2(num_classes=1, **kwargs),
             [16, 24, 48, 88, 120, 208, 352],
@@ -766,7 +782,7 @@ def __init__(self, layers=[1, 2, 4, 6], **kwargs):
 
 
 class EfficientNetB3Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b3(num_classes=1, **kwargs),
             [24, 32, 48, 96, 136, 232, 384],
@@ -776,7 +792,7 @@ def __init__(self, layers=[1, 2, 4, 6], **kwargs):
 
 
 class EfficientNetB4Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b4(num_classes=1, **kwargs),
             [24, 32, 56, 112, 160, 272, 448],
@@ -786,7 +802,7 @@ def __init__(self, layers=[1, 2, 4, 6], **kwargs):
 
 
 class EfficientNetB5Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b5(num_classes=1, **kwargs),
             [24, 40, 64, 128, 176, 304, 512],
@@ -796,7 +812,7 @@ def __init__(self, layers=[1, 2, 4, 6], **kwargs):
 
 
 class EfficientNetB6Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b6(num_classes=1, **kwargs),
             [32, 40, 72, 144, 200, 344, 576],
@@ -806,7 +822,7 @@ def __init__(self, layers=[1, 2, 4, 6], **kwargs):
 
 
 class EfficientNetB7Encoder(EfficientNetEncoder):
-    def __init__(self, layers=[1, 2, 4, 6], **kwargs):
+    def __init__(self, layers=None, **kwargs):
         super().__init__(
             efficient_net_b7(num_classes=1, **kwargs),
             [32, 48, 80, 160, 224, 384, 640],

tests/test_losses.py

@@ -9,8 +9,8 @@ def test_sigmoid_focal_loss():
     input_bad = torch.Tensor([-1, 2, 0]).float()
     target = torch.Tensor([1, 0, 1])
 
-    loss_good = F.sigmoid_focal_loss(input_good, target)
-    loss_bad = F.sigmoid_focal_loss(input_bad, target)
+    loss_good = F.focal_loss_with_logits(input_good, target)
+    loss_bad = F.focal_loss_with_logits(input_bad, target)
     assert loss_good < loss_bad