|
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
-
- import paddle
- from paddle import nn
- import paddle.nn.functional as F
-
- from paddleseg.cvlibs import manager
-
-
- @manager.LOSSES.add_component
- class DiceLoss(nn.Layer):
- """
- The implements of the dice loss.
-
- Args:
- weight (list[float], optional): The weight for each class. Default: None.
- ignore_index (int64): ignore_index (int64, optional): Specifies a target value that
- is ignored and does not contribute to the input gradient. Default ``255``.
- smooth (float32): Laplace smoothing to smooth dice loss and accelerate convergence.
- Default: 1.0
- """
-
- def __init__(self, weight=None, ignore_index=255, smooth=1.0):
- super().__init__()
- self.weight = weight
- self.ignore_index = ignore_index
- self.smooth = smooth
- self.eps = 1e-8
-
- def forward(self, logits, labels):
- num_class = logits.shape[1]
- if self.weight is not None:
- assert num_class == len(self.weight), \
- "The lenght of weight should be euqal to the num class"
-
- mask = labels != self.ignore_index
- mask = paddle.cast(paddle.unsqueeze(mask, 1), 'float32')
-
- labels[labels == self.ignore_index] = 0
- labels_one_hot = F.one_hot(labels, num_class)
- labels_one_hot = paddle.transpose(labels_one_hot, [0, 3, 1, 2])
- logits = F.softmax(logits, axis=1)
-
- dice_loss = 0.0
- for i in range(num_class):
- dice_loss_i = dice_loss_helper(logits[:, i], labels_one_hot[:, i],
- mask, self.smooth, self.eps)
- if self.weight is not None:
- dice_loss_i *= self.weight[i]
- dice_loss += dice_loss_i
- dice_loss = dice_loss / num_class
-
- return dice_loss
-
-
- def dice_loss_helper(logit, label, mask, smooth, eps):
- assert logit.shape == label.shape, \
- "The shape of logit and label should be the same"
- logit = paddle.reshape(logit, [0, -1])
- label = paddle.reshape(label, [0, -1])
- mask = paddle.reshape(mask, [0, -1])
- logit *= mask
- label *= mask
- intersection = paddle.sum(logit * label, axis=1)
- cardinality = paddle.sum(logit + label, axis=1)
- dice_loss = 1 - (2 * intersection + smooth) / (cardinality + smooth + eps)
- dice_loss = dice_loss.mean()
- return dice_loss
|