Merge pull request #39 from BloodAxe/develop

Release 0.3.1

Author: BloodAxe

pytorch_toolbelt/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import absolute_import
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"

pytorch_toolbelt/inference/tta.py

@@ -25,7 +25,7 @@
 
 def fliplr_image2label(model: nn.Module, image: Tensor) -> Tensor:
     """Test-time augmentation for image classification that averages predictions
-    for input image and vertically flipped one.
+    for input image and horizontally flipped one.
 
     :param model:
     :param image:
@@ -176,7 +176,7 @@ def d4_image2label(model: nn.Module, image: Tensor) -> Tensor:
 
 
 def d4_image2mask(model: nn.Module, image: Tensor) -> Tensor:
-    """Test-time augmentation for image classification that averages predictions
+    """Test-time augmentation for image segmentation that averages predictions
     of all D4 augmentations applied to input image.
 
     For segmentation we need to reverse the augmentation after making a prediction
@@ -222,7 +222,7 @@ class MultiscaleTTAWrapper(nn.Module):
     Multiscale TTA wrapper module
     """
 
-    def __init__(self, model: nn.Module, scale_levels: List[float]):
+    def __init__(self, model: nn.Module, scale_levels: List[float] = None, size_offsets: List[int] = None):
         """
         Initialize multi-scale TTA wrapper
 
@@ -231,9 +231,11 @@ def __init__(self, model: nn.Module, scale_levels: List[float]):
             e.g: [0.5, 0.75, 1.25]
         """
         super().__init__()
-        assert len(scale_levels)
+        assert scale_levels or size_offsets, "Either scale_levels or size_offsets must be set"
+        assert not (scale_levels and size_offsets), "Either scale_levels or size_offsets must be set"
         self.model = model
         self.scale_levels = scale_levels
+        self.size_offsets = size_offsets
 
     def forward(self, input: Tensor) -> Tensor:
         h = input.size(2)
@@ -242,11 +244,21 @@ def forward(self, input: Tensor) -> Tensor:
         out_size = h, w
         output = self.model(input)
 
-        for scale in self.scale_levels:
-            dst_size = int(h * scale), int(w * scale)
-            input_scaled = interpolate(input, dst_size, mode="bilinear", align_corners=False)
-            output_scaled = self.model(input_scaled)
-            output_scaled = interpolate(output_scaled, out_size, mode="bilinear", align_corners=False)
-            output += output_scaled
-
-        return output / (1 + len(self.scale_levels))
+        if self.scale_levels:
+            for scale in self.scale_levels:
+                dst_size = int(h * scale), int(w * scale)
+                input_scaled = interpolate(input, dst_size, mode="bilinear", align_corners=False)
+                output_scaled = self.model(input_scaled)
+                output_scaled = interpolate(output_scaled, out_size, mode="bilinear", align_corners=False)
+                output += output_scaled
+            output /= 1.0 + len(self.scale_levels)
+        elif self.size_offsets:
+            for offset in self.size_offsets:
+                dst_size = int(h + offset), int(w + offset)
+                input_scaled = interpolate(input, dst_size, mode="bilinear", align_corners=False)
+                output_scaled = self.model(input_scaled)
+                output_scaled = interpolate(output_scaled, out_size, mode="bilinear", align_corners=False)
+                output += output_scaled
+            output /= 1.0 + len(self.size_offsets)
+
+        return output

pytorch_toolbelt/losses/soft_ce.py

@@ -6,7 +6,7 @@
 
 
 class SoftCrossEntropyLoss(nn.Module):
-    def __init__(self, smooth_factor=1e-4, ignore_index=None):
+    def __init__(self, smooth_factor=1e-4, ignore_index: int = -100):
         super().__init__()
         self.smooth_factor = smooth_factor
         self.ignore_index = ignore_index

pytorch_toolbelt/modules/decoders/hrnet.py

@@ -1,19 +1,29 @@
 from collections import OrderedDict
 from typing import List
 
+import torch
+import torch.nn.functional as F
 from torch import nn, Tensor
 
-from .common import DecoderModule
+from .common import SegmentationDecoderModule
 
-__all__ = ["HRNetDecoder"]
+__all__ = ["HRNetSegmentationDecoder"]
 
 
-class HRNetDecoder(DecoderModule):
-    def __init__(self, feature_maps: List[int], output_channels: int, dropout=0.0):
+class HRNetSegmentationDecoder(SegmentationDecoderModule):
+    def __init__(
+        self,
+        feature_maps: List[int],
+        output_channels: int,
+        dropout=0.0,
+        interpolation_mode="nearest",
+        align_corners=None,
+    ):
         super().__init__()
+        self.interpolation_mode = interpolation_mode
+        self.align_corners = align_corners
 
-        features = feature_maps[-1]
-
+        features = sum(feature_maps)
         self.embedding = nn.Sequential(
             OrderedDict(
                 [
@@ -37,5 +47,15 @@ def __init__(self, feature_maps: List[int], output_channels: int, dropout=0.0):
         )
 
     def forward(self, features: List[Tensor]):
-        embedding = self.embedding(features[-1])
+        x_size = features[0].size()[2:]
+
+        resized_feature_maps = [features[0]]
+        for feature_map in features[1:]:
+            feature_map = F.interpolate(
+                feature_map, size=x_size, mode=self.interpolation_mode, align_corners=self.align_corners
+            )
+            resized_feature_maps.append(feature_map)
+
+        feature_map = torch.cat(resized_feature_maps, dim=1)
+        embedding = self.embedding(feature_map)
         return self.logits(embedding)

pytorch_toolbelt/modules/encoders/unet.py

@@ -16,16 +16,15 @@ def __init__(self, input_channels=3, features=32, num_layers=4, growth_factor=2,
         super().__init__(feature_maps, strides, layers=list(range(num_layers)))
 
         input_filters = input_channels
-        output_filters = feature_maps[0]
         self.num_layers = num_layers
         for layer in range(num_layers):
-            block = UnetEncoderBlock(input_filters, output_filters, abn_block=abn_block)
-
+            block = UnetEncoderBlock(input_filters, feature_maps[layer], abn_block=abn_block)
+            input_filters = feature_maps[layer]
             self.add_module(f"layer{layer}", block)
 
     @property
     def encoder_layers(self):
-        return [self[f"layer{layer}"] for layer in range(self.num_layers)]
+        return [self.__getattr__(f"layer{layer}") for layer in range(self.num_layers)]
 
     def change_input_channels(self, input_channels: int, mode="auto"):
         self.layer0.conv1 = make_n_channel_input(self.layer0.conv1, input_channels, mode)

pytorch_toolbelt/utils/catalyst/__init__.py

@@ -4,3 +4,4 @@
 from .metrics import *
 from .opl import *
 from .visualization import *
+from .utils import *

pytorch_toolbelt/utils/catalyst/metrics.py

@@ -227,7 +227,7 @@ def binary_dice_iou_score(
     if mode == "dice":
         score = (2.0 * intersection) / (cardinality + eps)
     else:
-        score = intersection / (cardinality + eps)
+        score = intersection / (cardinality - intersection + eps)
 
     has_targets = torch.sum(y_true) > 0
     has_predicted = torch.sum(y_pred) > 0

pytorch_toolbelt/utils/catalyst/utils.py

@@ -0,0 +1,45 @@
+from typing import Dict
+
+import torch
+
+__all__ = ["clean_checkpoint", "report_checkpoint"]
+
+
+def clean_checkpoint(src_fname, dst_fname):
+    """
+    Remove optimizer, scheduler and criterion states from checkpoint
+    :param src_fname: Source checkpoint filename
+    :param dst_fname: Target checkpoint filename (can be same)
+    """
+    checkpoint = torch.load(src_fname, map_location="cpu")
+
+    keys = ["criterion_state_dict", "optimizer_state_dict", "scheduler_state_dict"]
+
+    for key in keys:
+        if key in checkpoint:
+            del checkpoint[key]
+
+    torch.save(checkpoint, dst_fname)
+
+
+def report_checkpoint(checkpoint: Dict):
+    """
+    Print checkpoint metrics and epoch number
+    :param checkpoint:
+    """
+    print("Epoch          :", checkpoint["epoch"])
+
+    skip_fields = [
+        "_base/lr",
+        "_base/momentum",
+        "_timers/data_time",
+        "_timers/model_time",
+        "_timers/batch_time",
+        "_timers/_fps",
+    ]
+    print(
+        "Metrics (Train):", [(k, v) for k, v, in checkpoint["epoch_metrics"]["train"].items() if k not in skip_fields]
+    )
+    print(
+        "Metrics (Valid):", [(k, v) for k, v, in checkpoint["epoch_metrics"]["valid"].items() if k not in skip_fields]
+    )