Merge pull request #46 from BloodAxe/feature/new-catalyst

Feature/new catalyst

Author: BloodAxe

pytorch_toolbelt/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import absolute_import
 
-__version__ = "0.3.3"
+__version__ = "0.4.0"

pytorch_toolbelt/inference/ensembling.py

@@ -86,7 +86,7 @@ def forward(self, *input, **kwargs):  # skipcq: PYL-W0221
 
         if self.average:
             for key in keys:
-                output_0[key].mul_(1. / num_models)
+                output_0[key].mul_(1.0 / num_models)
 
         return output_0
 

pytorch_toolbelt/losses/dice.py

@@ -82,12 +82,12 @@ def forward(self, y_pred: Tensor, y_true: Tensor) -> Tensor:
             y_true = y_true.view(bs, num_classes, -1)
             y_pred = y_pred.view(bs, num_classes, -1)
 
-        scores = soft_dice_score(y_pred, y_true.type_as(y_pred), self.smooth, self.eps, dims=dims)
+        scores = soft_dice_score(y_pred, y_true.type_as(y_pred), smooth=self.smooth, eps=self.eps, dims=dims)
 
         if self.log_loss:
             loss = -torch.log(scores.clamp_min(self.eps))
         else:
-            loss = 1 - scores
+            loss = 1.0 - scores
 
         # Dice loss is undefined for non-empty classes
         # So we zero contribution of channel that does not have true pixels

pytorch_toolbelt/losses/functional.py

@@ -8,7 +8,7 @@
 
 
 def focal_loss_with_logits(
-    input: torch.Tensor,
+    output: torch.Tensor,
     target: torch.Tensor,
     gamma: float = 2.0,
     alpha: Optional[float] = 0.25,
@@ -22,7 +22,7 @@ def focal_loss_with_logits(
     See :class:`~pytorch_toolbelt.losses.FocalLoss` for details.
 
     Args:
-        input: Tensor of arbitrary shape
+        output: Tensor of arbitrary shape (predictions of the model)
         target: Tensor of the same shape as input
         gamma: Focal loss power factor
         alpha: Weight factor to balance positive and negative samples. Alpha must be in [0...1] range,
@@ -40,9 +40,9 @@ def focal_loss_with_logits(
     References:
         https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/loss/losses.py
     """
-    target = target.type(input.type())
+    target = target.type(output.type())
 
-    logpt = F.binary_cross_entropy_with_logits(input, target, reduction="none")
+    logpt = F.binary_cross_entropy_with_logits(output, target, reduction="none")
     pt = torch.exp(-logpt)
 
     # compute the loss
@@ -76,19 +76,22 @@ def 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"):
+def reduced_focal_loss(output: torch.Tensor, target: torch.Tensor, threshold=0.5, gamma=2.0, reduction="mean"):
     return focal_loss_with_logits(
-        input, target, alpha=None, gamma=gamma, reduction=reduction, reduced_threshold=threshold
+        output, target, alpha=None, gamma=gamma, reduction=reduction, reduced_threshold=threshold
     )
 
 
-def soft_jaccard_score(y_pred: torch.Tensor, y_true: torch.Tensor, smooth=0.0, eps=1e-7, dims=None) -> torch.Tensor:
+def soft_jaccard_score(
+    output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
+) -> torch.Tensor:
     """
 
-    :param y_pred:
-    :param y_true:
+    :param output:
+    :param target:
     :param smooth:
     :param eps:
+    :param dims:
     :return:
 
     Shape:
@@ -98,25 +101,27 @@ def soft_jaccard_score(y_pred: torch.Tensor, y_true: torch.Tensor, smooth=0.0, e
         - Output: scalar.
 
     """
-    assert y_pred.size() == y_true.size()
+    assert output.size() == target.size()
 
     if dims is not None:
-        intersection = torch.sum(y_pred * y_true, dim=dims)
-        cardinality = torch.sum(y_pred + y_true, dim=dims)
+        intersection = torch.sum(output * target, dim=dims)
+        cardinality = torch.sum(output + target, dim=dims)
     else:
-        intersection = torch.sum(y_pred * y_true)
-        cardinality = torch.sum(y_pred + y_true)
+        intersection = torch.sum(output * target)
+        cardinality = torch.sum(output + target)
 
     union = cardinality - intersection
-    jaccard_score = (intersection + smooth) / (union.clamp_min(eps) + smooth)
+    jaccard_score = (intersection + smooth) / (union + smooth).clamp_min(eps)
     return jaccard_score
 
 
-def soft_dice_score(y_pred: torch.Tensor, y_true: torch.Tensor, smooth=0, eps=1e-7, dims=None) -> torch.Tensor:
+def soft_dice_score(
+    output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
+) -> torch.Tensor:
     """
 
-    :param y_pred:
-    :param y_true:
+    :param output:
+    :param target:
     :param smooth:
     :param eps:
     :return:
@@ -128,28 +133,28 @@ def soft_dice_score(y_pred: torch.Tensor, y_true: torch.Tensor, smooth=0, eps=1e
         - Output: scalar.
 
     """
-    assert y_pred.size() == y_true.size()
+    assert output.size() == target.size()
     if dims is not None:
-        intersection = torch.sum(y_pred * y_true, dim=dims)
-        cardinality = torch.sum(y_pred + y_true, dim=dims)
+        intersection = torch.sum(output * target, dim=dims)
+        cardinality = torch.sum(output + target, dim=dims)
     else:
-        intersection = torch.sum(y_pred * y_true)
-        cardinality = torch.sum(y_pred + y_true)
-    dice_score = (2.0 * intersection + smooth) / (cardinality.clamp_min(eps) + smooth)
+        intersection = torch.sum(output * target)
+        cardinality = torch.sum(output + target)
+    dice_score = (2.0 * intersection + smooth) / (cardinality + smooth).clamp_min(eps)
     return dice_score
 
 
-def wing_loss(prediction: torch.Tensor, target: torch.Tensor, width=5, curvature=0.5, reduction="mean"):
+def wing_loss(output: torch.Tensor, target: torch.Tensor, width=5, curvature=0.5, reduction="mean"):
     """
     https://arxiv.org/pdf/1711.06753.pdf
-    :param prediction:
+    :param output:
     :param target:
     :param width:
     :param curvature:
     :param reduction:
     :return:
     """
-    diff_abs = (target - prediction).abs()
+    diff_abs = (target - output).abs()
     loss = diff_abs.clone()
 
     idx_smaller = diff_abs < width
@@ -180,7 +185,7 @@ def label_smoothed_nll_loss(
     :param target:
     :param epsilon:
     :param ignore_index:
-    :param reduce:
+    :param reduction:
     :return:
     """
     if target.dim() == lprobs.dim() - 1:

pytorch_toolbelt/losses/jaccard.py

@@ -82,12 +82,12 @@ def forward(self, y_pred: Tensor, y_true: Tensor) -> Tensor:
             y_true = y_true.view(bs, num_classes, -1)
             y_pred = y_pred.view(bs, num_classes, -1)
 
-        scores = soft_jaccard_score(y_pred, y_true.type(y_pred.dtype), self.smooth, self.eps, dims=dims)
+        scores = soft_jaccard_score(y_pred, y_true.type(y_pred.dtype), smooth=self.smooth, eps=self.eps, dims=dims)
 
         if self.log_loss:
             loss = -torch.log(scores.clamp_min(self.eps))
         else:
-            loss = 1 - scores
+            loss = 1.0 - scores
 
         # IoU loss is defined for non-empty classes
         # So we zero contribution of channel that does not have true pixels

pytorch_toolbelt/modules/activations.py

@@ -129,7 +129,8 @@ def mish_jit_bwd(x, grad_output):
 
 
 class MishFunction(torch.autograd.Function):
-    """ Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """
+    Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
     A memory efficient, jit scripted variant of Mish
     """
 

pytorch_toolbelt/utils/catalyst/criterions.py

@@ -1,8 +1,7 @@
 import math
 
 import torch
-from catalyst.dl import CriterionCallback, RunnerState
-from catalyst.dl.callbacks.criterion import _add_loss_to_state
+from catalyst.dl import IRunner, CriterionCallback
 from torch import nn
 from torch.nn import functional as F
 
@@ -67,22 +66,22 @@ def __init__(
         self.p = p
         self.multiplier = None
 
-    def on_loader_start(self, state: RunnerState):
-        self.is_needed = not self.on_train_only or state.loader_name.startswith("train")
+    def on_loader_start(self, runner: IRunner):
+        self.is_needed = not self.on_train_only or runner.loader_name.startswith("train")
         if self.is_needed:
-            state.metrics.epoch_values[state.loader_name][f"l{self.p}_weight_decay"] = self.multiplier
+            runner.metrics.epoch_values[runner.loader_name][f"l{self.p}_weight_decay"] = self.multiplier
 
-    def on_epoch_start(self, state: RunnerState):
-        training_progress = float(state.epoch) / float(state.num_epochs)
+    def on_epoch_start(self, runner: IRunner):
+        training_progress = float(runner.epoch) / float(runner.num_epochs)
         self.multiplier = get_multiplier(training_progress, self.schedule, self.start_wd, self.end_wd)
 
-    def on_batch_end(self, state: RunnerState):
+    def on_batch_end(self, runner: IRunner):
         if not self.is_needed:
             return
 
         lp_reg = 0
 
-        for module in state.model.children():
+        for module in runner.model.children():
             if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm1d, nn.BatchNorm3d)):
                 continue
 
@@ -93,8 +92,7 @@ def on_batch_end(self, state: RunnerState):
                 if param.requires_grad:
                     lp_reg = param.norm(self.p) * self.multiplier + lp_reg
 
-        state.metrics.add_batch_value(metrics_dict={self.prefix: lp_reg.item()})
-        _add_loss_to_state(self.prefix, state, lp_reg)
+        runner.batch_metrics.update(**{self.prefix: lp_reg.item()})
 
 
 class TSACriterionCallback(CriterionCallback):
@@ -117,7 +115,7 @@ def __init__(
         prefix: str = "loss",
         criterion_key: str = None,
         multiplier: float = 1.0,
-        unsupervised_label=-100,
+        ignore_index=-100,
     ):
         super().__init__(
             input_key=input_key,
@@ -129,7 +127,7 @@ def __init__(
         self.num_epochs = num_epochs
         self.num_classes = num_classes
         self.tsa_threshold = None
-        self.unsupervised_label = unsupervised_label
+        self.ignore_index = ignore_index
 
     def get_tsa_threshold(self, current_epoch, schedule, start, end) -> float:
         training_progress = float(current_epoch) / float(self.num_epochs)
@@ -148,16 +146,16 @@ def get_tsa_threshold(self, current_epoch, schedule, start, end) -> float:
             raise KeyError(f"Unsupported schedule name {schedule}")
         return threshold * (end - start) + start
 
-    def on_epoch_start(self, state: RunnerState):
-        if state.loader_name == "train":
-            self.tsa_threshold = self.get_tsa_threshold(state.epoch, "exp_schedule", 1.0 / self.num_classes, 1.0)
-            state.metrics.epoch_values["train"]["tsa_threshold"] = self.tsa_threshold
+    def on_epoch_start(self, runner: IRunner):
+        if runner.loader_name == "train":
+            self.tsa_threshold = self.get_tsa_threshold(runner.epoch, "exp_schedule", 1.0 / self.num_classes, 1.0)
+            runner.epoch_metrics["train"]["tsa_threshold"] = self.tsa_threshold
 
-    def _compute_loss(self, state: RunnerState, criterion):
+    def _compute_loss(self, runner: IRunner, criterion):
 
-        logits = state.output[self.output_key]
-        targets = state.input[self.input_key]
-        supervised_mask = targets != self.unsupervised_label  # Mask indicating labeled samples
+        logits = runner.output[self.output_key]
+        targets = runner.input[self.input_key]
+        supervised_mask = targets != self.ignore_index  # Mask indicating labeled samples
 
         targets = targets[supervised_mask]
         logits = logits[supervised_mask]

pytorch_toolbelt/utils/catalyst/loss_adapter.py

@@ -1,21 +1,15 @@
+from typing import Dict
+
 import torch
 from catalyst.dl import (
     CriterionCallback,
-    RunnerState,
-    OptimizerCallback,
-    CheckpointCallback,
-    SchedulerCallback,
-    SupervisedExperiment,
-    Callback,
+    IRunner,
 )
-from catalyst.dl.callbacks import VerboseLogger, ConsoleLogger, TensorboardLogger, RaiseExceptionCallback
 from torch import nn, Tensor
-from typing import Dict
 
 __all__ = [
     "TrainOnlyCriterionCallback",
     "PassthroughCriterionCallback",
-    "ParallelLossSupervisedExperiment",
     "LossModule",
     "LossWrapper",
 ]
@@ -26,7 +20,7 @@ class TrainOnlyCriterionCallback(CriterionCallback):
     Computes loss only on training stage
     """
 
-    def _compute_loss_value(self, state: RunnerState, criterion):
+    def _compute_loss_value(self, state: IRunner, criterion):
         predictions = self._get_output(state.output, self.output_key)
         targets = self._get_input(state.input, self.input_key)
 
@@ -36,7 +30,7 @@ def _compute_loss_value(self, state: RunnerState, criterion):
         loss = criterion(predictions, targets)
         return loss
 
-    def _compute_loss_key_value(self, state: RunnerState, criterion):
+    def _compute_loss_key_value(self, state: IRunner, criterion):
         output = self._get_output(state.output, self.output_key)
         input = self._get_input(state.input, self.input_key)
 
@@ -47,44 +41,6 @@ def _compute_loss_key_value(self, state: RunnerState, criterion):
         return loss
 
 
-class ParallelLossSupervisedExperiment(SupervisedExperiment):
-    """
-    Custom experiment class. To use in conjunction with LossWrapper.
-    """
-
-    def get_callbacks(self, stage: str) -> "OrderedDict[str, Callback]":
-        """
-        Override of ``BaseExperiment.get_callbacks`` method.
-        Will add several of callbacks by default in case they missed.
-
-        Args:
-            stage (str): name of stage. It should start with `infer` if you
-                don't need default callbacks, as they required only for
-                training stages.
-        Returns:
-            List[Callback]: list of callbacks for experiment
-        """
-        callbacks = self._callbacks
-        default_callbacks = []
-        if self._verbose:
-            default_callbacks.append(("verbose", VerboseLogger))
-        if not stage.startswith("infer"):
-            # default_callbacks.append(("_criterion", CriterionCallback)) # Commented
-            default_callbacks.append(("_optimizer", OptimizerCallback))
-            if self._scheduler is not None:
-                default_callbacks.append(("_scheduler", SchedulerCallback))
-            default_callbacks.append(("_saver", CheckpointCallback))
-            default_callbacks.append(("console", ConsoleLogger))
-            default_callbacks.append(("tensorboard", TensorboardLogger))
-        default_callbacks.append(("exception", RaiseExceptionCallback))
-
-        for callback_name, callback_fn in default_callbacks:
-            is_already_present = any(isinstance(x, callback_fn) for x in callbacks.values())
-            if not is_already_present:
-                callbacks[callback_name] = callback_fn()
-        return callbacks
-
-
 class PassthroughCriterionCallback(CriterionCallback):
     """
     Returns one of model's outputs as loss values
@@ -93,11 +49,11 @@ class PassthroughCriterionCallback(CriterionCallback):
     def __init__(self, output_key, multiplier=1.0):
         super().__init__(output_key=output_key, prefix=output_key, multiplier=multiplier)
 
-    def _compute_loss_value(self, state: RunnerState, criterion):
+    def _compute_loss_value(self, state: IRunner, criterion):
         loss = self._get_output(state.output, self.output_key)
         return loss.mean()
 
-    def _compute_loss_key_value(self, state: RunnerState, criterion):
+    def _compute_loss_key_value(self, state: IRunner, criterion):
         loss = self._get_output(state.output, self.output_key)
         return loss.mean()