Merge pull request #14 from BloodAxe/develop

Release 0.1.0

Author: BloodAxe

demo/demo_losses.py

@@ -9,33 +9,50 @@
 def main():
     losses = {
         "bce": BCEWithLogitsLoss(),
-        "focal": L.BinaryFocalLoss(),
-        "jaccard": L.BinaryJaccardLoss(),
+        # "focal": L.BinaryFocalLoss(),
+        # "jaccard": L.BinaryJaccardLoss(),
         # "jaccard_log": L.BinaryJaccardLogLoss(),
-        "lovasz": L.BinaryLovaszLoss(),
-        # "bce+jaccard_log": L.BinaryJaccardLogLoss(),
-        "reduced_focal": L.BinaryFocalLoss(reduced=True)
+        # "dice": L.BinaryDiceLoss(),
+        # "dice_log": L.BinaryDiceLogLoss(),
+        # "sdice": L.BinarySymmetricDiceLoss(),
+        # "sdice_log": L.BinarySymmetricDiceLoss(log_loss=True),
+
+        "bce+lovasz": L.JointLoss(BCEWithLogitsLoss(), L.BinaryLovaszLoss()),
+        # "lovasz": L.BinaryLovaszLoss(),
+        # "bce+jaccard": L.JointLoss(BCEWithLogitsLoss(),
+        #                            L.BinaryJaccardLoss(), 1, 0.5),
+
+        # "bce+log_jaccard": L.JointLoss(BCEWithLogitsLoss(),
+        #                            L.BinaryJaccardLogLoss(), 1, 0.5),
+
+        # "bce+log_dice": L.JointLoss(BCEWithLogitsLoss(),
+        #                                L.BinaryDiceLogLoss(), 1, 0.5)
+
+        # "reduced_focal": L.BinaryFocalLoss(reduced=True)
     }
 
-    x_vec = torch.arange(-5, 5, 0.01)
+    dx = 0.01
+    x_vec = torch.arange(-5, 5, dx).view(-1, 1).expand((-1, 100))
 
-    plt.figure()
+    f, ax = plt.subplots(3, figsize=(16, 16))
 
     for name, loss in losses.items():
         x_arr = []
         y_arr = []
-        target = torch.tensor(1.0)
+        target = torch.tensor(1.0).view(1).expand((100))
 
         for x in x_vec:
-            y = loss(x, target)
+            y = loss(x, target).item()
 
-            x_arr.append(float(x))
+            x_arr.append(float(x[0]))
             y_arr.append(float(y))
 
-        plt.plot(x_arr, y_arr, label=name)
+        ax[0].plot(x_arr, y_arr, label=name)
+        ax[1].plot(x_arr, np.gradient(y_arr, dx))
+        ax[2].plot(x_arr, np.gradient(np.gradient(y_arr, dx), dx))
 
-    plt.legend()
-    plt.show()
+    f.legend()
+    f.show()
 
 
 if __name__ == '__main__':

pytorch_toolbelt/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import absolute_import
 
-__version__ = '0.0.9'
+__version__ = '0.1.0'

pytorch_toolbelt/inference/tta.py

@@ -4,13 +4,20 @@
 transformation written in PyTorch and respect gradients flow.
 """
 from functools import partial
-from typing import Tuple
+from typing import Tuple, List
 
 from torch import Tensor, nn
+
 from . import functional as F
 
-__all__ = ['d4_image2label', 'd4_image2mask', 'fivecrop_image2label', 'fliplr_image2mask',
-           'fliplr_image2label', 'TTAWrapper']
+__all__ = ['d4_image2label',
+           'd4_image2mask',
+           'fivecrop_image2label',
+           'tencrop_image2label',
+           'fliplr_image2mask',
+           'fliplr_image2label',
+           'TTAWrapper',
+           'MultiscaleTTAWrapper']
 
 
 def fliplr_image2label(model: nn.Module, image: Tensor) -> Tensor:
@@ -26,7 +33,8 @@ def fliplr_image2label(model: nn.Module, image: Tensor) -> Tensor:
     return output * one_over_2
 
 
-def fivecrop_image2label(model: nn.Module, image: Tensor, crop_size: Tuple) -> Tensor:
+def fivecrop_image2label(model: nn.Module, image: Tensor,
+                         crop_size: Tuple) -> Tensor:
     """Test-time augmentation for image classification that takes five crops out of input tensor (4 on corners and central)
     and averages predictions from them.
 
@@ -61,16 +69,19 @@ def fivecrop_image2label(model: nn.Module, image: Tensor, crop_size: Tuple) -> T
     center_crop_y = (image_height - crop_height) // 2
     center_crop_x = (image_width - crop_width) // 2
 
-    crop_cc = image[..., center_crop_y:center_crop_y + crop_height, center_crop_x:center_crop_x + crop_width]
+    crop_cc = image[..., center_crop_y:center_crop_y + crop_height,
+                         center_crop_x:center_crop_x + crop_width]
     assert crop_cc.size(2) == crop_height
     assert crop_cc.size(3) == crop_width
 
-    output = model(crop_tl) + model(crop_tr) + model(crop_bl) + model(crop_br) + model(crop_cc)
+    output = model(crop_tl) + model(crop_tr) + model(crop_bl) + model(
+        crop_br) + model(crop_cc)
     one_over_5 = float(1.0 / 5.0)
     return output * one_over_5
 
 
-def tencrop_image2label(model: nn.Module, image: Tensor, crop_size: Tuple) -> Tensor:
+def tencrop_image2label(model: nn.Module, image: Tensor,
+                        crop_size: Tuple) -> Tensor:
     """Test-time augmentation for image classification that takes five crops out of input tensor (4 on corners and central)
     and averages predictions from them and from their horisontally-flipped versions (10-Crop TTA).
 
@@ -105,7 +116,8 @@ def tencrop_image2label(model: nn.Module, image: Tensor, crop_size: Tuple) -> Te
     center_crop_y = (image_height - crop_height) // 2
     center_crop_x = (image_width - crop_width) // 2
 
-    crop_cc = image[..., center_crop_y:center_crop_y + crop_height, center_crop_x:center_crop_x + crop_width]
+    crop_cc = image[..., center_crop_y:center_crop_y + crop_height,
+                         center_crop_x:center_crop_x + crop_width]
     assert crop_cc.size(2) == crop_height
     assert crop_cc.size(3) == crop_width
 
@@ -170,13 +182,16 @@ def d4_image2mask(model: nn.Module, image: Tensor) -> Tensor:
     """
     output = model(image)
 
-    for aug, deaug in zip([F.torch_rot90, F.torch_rot180, F.torch_rot270], [F.torch_rot270, F.torch_rot180, F.torch_rot90]):
+    for aug, deaug in zip([F.torch_rot90, F.torch_rot180, F.torch_rot270],
+                          [F.torch_rot270, F.torch_rot180, F.torch_rot90]):
         x = deaug(model(aug(image)))
         output = output + x
 
     image = F.torch_transpose(image)
 
-    for aug, deaug in zip([F.torch_none, F.torch_rot90, F.torch_rot180, F.torch_rot270], [F.torch_none, F.torch_rot270, F.torch_rot180, F.torch_rot90]):
+    for aug, deaug in zip(
+            [F.torch_none, F.torch_rot90, F.torch_rot180, F.torch_rot270],
+            [F.torch_none, F.torch_rot270, F.torch_rot180, F.torch_rot90]):
         x = deaug(model(aug(image)))
         output = output + F.torch_transpose(x)
 
@@ -185,10 +200,47 @@ def d4_image2mask(model: nn.Module, image: Tensor) -> Tensor:
 
 
 class TTAWrapper(nn.Module):
-    def __init__(self, model, tta_function, **kwargs):
+    def __init__(self, model: nn.Module, tta_function, **kwargs):
         super().__init__()
         self.model = model
         self.tta = partial(tta_function, **kwargs)
 
     def forward(self, *input):
         return self.tta(self.model, *input)
+
+
+class MultiscaleTTAWrapper(nn.Module):
+    """
+    Multiscale TTA wrapper module
+    """
+
+    def __init__(self, model: nn.Module, scale_levels: List[float]):
+        """
+        Initialize multi-scale TTA wrapper
+
+        :param model: Base model for inference
+        :param scale_levels: List of additional scale levels,
+            e.g: [0.5, 0.75, 1.25]
+        """
+        super().__init__()
+        assert len(scale_levels)
+        self.model = model
+        self.scale_levels = scale_levels
+
+    def forward(self, input: Tensor) -> Tensor:
+        h = input.size(2)
+        w = input.size(3)
+
+        out_size = h, w
+        output = self.model(input)
+
+        for scale in self.scale_levels:
+            dst_size = int(h * scale), int(w * scale)
+            input_scaled = F.interpolate(input, dst_size, mode='bilinear',
+                                         align_corners=True)
+            output_scaled = self.model(input_scaled)
+            output_scaled = F.interpolate(output_scaled, out_size,
+                                          mode='bilinear', align_corners=True)
+            output += output_scaled
+
+        return output / (1 + len(self.scale_levels))

pytorch_toolbelt/modules/abn.py

@@ -1,23 +1,20 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as functional
+from pytorch_toolbelt.modules.activations import ACT_LEAKY_RELU, ACT_NONE, \
+    ACT_HARD_SIGMOID, ACT_HARD_SWISH, ACT_SWISH, ACT_SELU, ACT_ELU, ACT_RELU6, \
+    ACT_RELU, hard_swish, hard_sigmoid, swish
 
-__all__ = ['ACT_RELU', 'ACT_ELU', 'ACT_SELU', 'ACT_NONE', 'ACT_LEAKY_RELU', 'ABN']
-
-# Activation names
-ACT_RELU = "relu"
-ACT_LEAKY_RELU = "leaky_relu"
-ACT_ELU = "elu"
-ACT_NONE = "none"
-ACT_SELU = "selu"
+__all__ = ['ABN']
 
 
 class ABN(nn.Module):
     """Activated Batch Normalization
     This gathers a `BatchNorm2d` and an activation function in a single module
     """
 
-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
+                 activation="leaky_relu", slope=0.01):
         """Create an Activated Batch Normalization module
         Parameters
         ----------
@@ -52,26 +49,39 @@ def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation
         self.reset_parameters()
 
     def reset_parameters(self):
-        nn.init.constant_(self.running_mean, 0)
-        nn.init.constant_(self.running_var, 1)
+        nn.init.zeros_(self.running_mean)
+        nn.init.ones_(self.running_var)
         if self.affine:
-            nn.init.constant_(self.weight, 1)
-            nn.init.constant_(self.bias, 0)
+            nn.init.ones_(self.weight)
+            nn.init.zeros_(self.bias)
 
     def forward(self, x):
-        x = functional.batch_norm(x, self.running_mean, self.running_var, self.weight, self.bias,
+        x = functional.batch_norm(x,
+                                  self.running_mean, self.running_var,
+                                  self.weight, self.bias,
                                   self.training, self.momentum, self.eps)
 
         if self.activation == ACT_RELU:
             return functional.relu(x, inplace=True)
+        elif self.activation == ACT_RELU6:
+            return functional.relu6(x, inplace=True)
         elif self.activation == ACT_LEAKY_RELU:
-            return functional.leaky_relu(x, negative_slope=self.slope, inplace=True)
+            return functional.leaky_relu(x, negative_slope=self.slope,
+                                         inplace=True)
         elif self.activation == ACT_ELU:
             return functional.elu(x, inplace=True)
         elif self.activation == ACT_SELU:
             return functional.selu(x, inplace=True)
-        else:
+        elif self.activation == ACT_SWISH:
+            return swish(x)
+        elif self.activation == ACT_HARD_SWISH:
+            return hard_swish(x, inplace=True)
+        elif self.activation == ACT_HARD_SIGMOID:
+            return hard_sigmoid(x, inplace=True)
+        elif self.activation == ACT_NONE:
             return x
+        else:
+            raise KeyError(self.activation)
 
     def __repr__(self):
         rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \

pytorch_toolbelt/modules/activations.py

@@ -3,6 +3,24 @@
 from torch import nn
 from torch.nn import functional as F
 
+__all__ = ['ACT_ELU',
+           'ACT_HARD_SIGMOID', 'ACT_HARD_SWISH', 'ACT_LEAKY_RELU', 'ACT_NONE',
+           'ACT_RELU', 'ACT_RELU6', 'ACT_SELU', 'ACT_SWISH',
+           'swish', 'hard_sigmoid', 'hard_swish', 'HardSigmoid', 'HardSwish',
+           'Swish', 'get_activation_module'
+           ]
+
+# Activation names
+ACT_RELU = "relu"
+ACT_RELU6 = "relu6"
+ACT_LEAKY_RELU = "leaky_relu"
+ACT_ELU = "elu"
+ACT_NONE = "none"
+ACT_SELU = "selu"
+ACT_SWISH = "swish"
+ACT_HARD_SWISH = "hard_swish"
+ACT_HARD_SIGMOID = "hard_sigmoid"
+
 
 def swish(x):
     return x * x.sigmoid()
@@ -70,6 +88,9 @@ def get_activation_module(activation_name: str, **kwargs) -> nn.Module:
     if activation_name.lower() == 'hard_sigmoid':
         return partial(HardSigmoid, **kwargs)
 
+    if activation_name.lower() == 'swish':
+        return partial(Swish, **kwargs)
+
     if activation_name.lower() == 'hard_swish':
         return partial(HardSwish, **kwargs)