torch_em.loss.dice

  1import torch
  2import torch.nn as nn
  3
  4
  5# TODO refactor
  6def flatten_samples(input_):
  7    """
  8    Flattens a tensor or a variable such that the channel axis is first and the sample axis
  9    is second. The shapes are transformed as follows:
 10        (N, C, H, W) --> (C, N * H * W)
 11        (N, C, D, H, W) --> (C, N * D * H * W)
 12        (N, C) --> (C, N)
 13    The input must be atleast 2d.
 14    """
 15    # Get number of channels
 16    num_channels = input_.size(1)
 17    # Permute the channel axis to first
 18    permute_axes = list(range(input_.dim()))
 19    permute_axes[0], permute_axes[1] = permute_axes[1], permute_axes[0]
 20    # For input shape (say) NCHW, this should have the shape CNHW
 21    permuted = input_.permute(*permute_axes).contiguous()
 22    # Now flatten out all but the first axis and return
 23    flattened = permuted.view(num_channels, -1)
 24    return flattened
 25
 26
 27def dice_score(input_, target, invert=False, channelwise=True, reduce_channel="sum", eps=1e-7):
 28    if input_.shape != target.shape:
 29        raise ValueError(f"Expect input and target of same shape, got: {input_.shape}, {target.shape}.")
 30
 31    if channelwise:
 32        # Flatten input and target to have the shape (C, N),
 33        # where N is the number of samples
 34        input_ = flatten_samples(input_)
 35        target = flatten_samples(target)
 36        # Compute numerator and denominator (by summing over samples and
 37        # leaving the channels intact)
 38        numerator = (input_ * target).sum(-1)
 39        denominator = (input_ * input_).sum(-1) + (target * target).sum(-1)
 40        channelwise_score = 2 * (numerator / denominator.clamp(min=eps))
 41        if invert:
 42            channelwise_score = 1. - channelwise_score
 43
 44        # Reduce the dice score over the channels to compute the overall dice score.
 45        # (default is to use the sum)
 46        if reduce_channel is None:
 47            score = channelwise_score
 48        elif reduce_channel == "sum":
 49            score = channelwise_score.sum()
 50        elif reduce_channel == "mean":
 51            score = channelwise_score.mean()
 52        elif reduce_channel == "max":
 53            score = channelwise_score.max()
 54        elif reduce_channel == "min":
 55            score = channelwise_score.min()
 56        else:
 57            raise ValueError(f"Unsupported channel reduction {reduce_channel}")
 58
 59    else:
 60        numerator = (input_ * target).sum()
 61        denominator = (input_ * input_).sum() + (target * target).sum()
 62        score = 2. * (numerator / denominator.clamp(min=eps))
 63        if invert:
 64            score = 1. - score
 65
 66    return score
 67
 68
 69class DiceLoss(nn.Module):
 70    def __init__(self, channelwise=True, eps=1e-7, reduce_channel="sum"):
 71        if reduce_channel not in ("sum", "mean", "max", "min", None):
 72            raise ValueError(f"Unsupported channel reduction {reduce_channel}")
 73        super().__init__()
 74        self.channelwise = channelwise
 75        self.eps = eps
 76        self.reduce_channel = reduce_channel
 77
 78        # all torch_em classes should store init kwargs to easily recreate the init call
 79        self.init_kwargs = {"channelwise": channelwise, "eps": self.eps, "reduce_channel": self.reduce_channel}
 80
 81    def forward(self, input_, target):
 82        return dice_score(input_, target,
 83                          invert=True, channelwise=self.channelwise,
 84                          eps=self.eps, reduce_channel=self.reduce_channel)
 85
 86
 87class DiceLossWithLogits(nn.Module):
 88    def __init__(self, channelwise=True, eps=1e-7, reduce_channel="sum"):
 89        super().__init__()
 90        self.channelwise = channelwise
 91        self.eps = eps
 92        self.reduce_channel = reduce_channel
 93
 94        # all torch_em classes should store init kwargs to easily recreate the init call
 95        self.init_kwargs = {"channelwise": channelwise, "eps": self.eps, "reduce_channel": self.reduce_channel}
 96
 97    def forward(self, input_, target):
 98        return dice_score(
 99            nn.functional.sigmoid(input_),
100            target,
101            invert=True,
102            channelwise=self.channelwise,
103            eps=self.eps,
104            reduce_channel=self.reduce_channel,
105        )
106
107
108class BCEDiceLoss(nn.Module):
109
110    def __init__(self, alpha=1., beta=1., channelwise=True, eps=1e-7):
111        super().__init__()
112        self.alpha = alpha
113        self.beta = beta
114        self.channelwise = channelwise
115        self.eps = eps
116
117        # all torch_em classes should store init kwargs to easily recreate the init call
118        self.init_kwargs = {"alpha": alpha, "beta": beta, "channelwise": channelwise, "eps": self.eps}
119
120    def forward(self, input_, target):
121        loss_dice = dice_score(
122            input_,
123            target,
124            invert=True,
125            channelwise=self.channelwise,
126            eps=self.eps
127        )
128        loss_bce = nn.functional.binary_cross_entropy(
129            input_, target
130        )
131        return self.alpha * loss_dice + self.beta * loss_bce
132
133
134# TODO think about how to handle combined losses like this for mixed precision training
135class BCEDiceLossWithLogits(nn.Module):
136
137    def __init__(self, alpha=1., beta=1., channelwise=True, eps=1e-7):
138        super().__init__()
139        self.alpha = alpha
140        self.beta = beta
141        self.channelwise = channelwise
142        self.eps = eps
143
144        # all torch_em classes should store init kwargs to easily recreate the init call
145        self.init_kwargs = {"alpha": alpha, "beta": beta, "channelwise": channelwise, "eps": self.eps}
146
147    def forward(self, input_, target):
148        loss_dice = dice_score(
149            torch.sigmoid(input_),
150            target,
151            invert=True,
152            channelwise=self.channelwise,
153            eps=self.eps
154        )
155        loss_bce = nn.functional.binary_cross_entropy_with_logits(
156            input_, target
157        )
158        return self.alpha * loss_dice + self.beta * loss_bce