decoupledsegnet_relax_boundary_loss.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import numpy as np
  15. import paddle
  16. from paddle import nn
  17. import paddle.nn.functional as F
  18. from scipy.ndimage.interpolation import shift
  19. from paddlex.paddleseg.cvlibs import manager
  20. @manager.LOSSES.add_component
  21. class RelaxBoundaryLoss(nn.Layer):
  22. """
  23. Implements the ohem cross entropy loss function.
  24. Args:
  25. border (int, optional): The value of border to relax. Default: 1.
  26. calculate_weights (bool, optional): Whether to calculate weights for every classes. Default: False.
  27. upper_bound (float, optional): The upper bound of weights if calculating weights for every classes. Default: 1.0.
  28. ignore_index (int64): Specifies a target value that is ignored
  29. and does not contribute to the input gradient. Default: 255.
  30. """
  31. def __init__(self,
  32. border=1,
  33. calculate_weights=False,
  34. upper_bound=1.0,
  35. ignore_index=255):
  36. super(RelaxBoundaryLoss, self).__init__()
  37. self.border = border
  38. self.calculate_weights = calculate_weights
  39. self.upper_bound = upper_bound
  40. self.ignore_index = ignore_index
  41. self.EPS = 1e-5
  42. def relax_onehot(self, label, num_classes):
  43. # pad label, and let ignore_index as num_classes
  44. if len(label.shape) == 3:
  45. label = label.unsqueeze(1)
  46. h, w = label.shape[-2], label.shape[-1]
  47. label = F.pad(label, [self.border] * 4, value=num_classes)
  48. label = label.squeeze(1)
  49. ignore_mask = (label == self.ignore_index).astype('int64')
  50. label = label * (1 - ignore_mask) + num_classes * ignore_mask
  51. onehot = 0
  52. for i in range(-self.border, self.border + 1):
  53. for j in range(-self.border, self.border + 1):
  54. h_start, h_end = 1 + i, h + 1 + i
  55. w_start, w_end = 1 + j, w + 1 + j
  56. label_ = label[:, h_start:h_end, w_start:w_end]
  57. onehot_ = F.one_hot(label_, num_classes + 1)
  58. onehot += onehot_
  59. onehot = (onehot > 0).astype('int64')
  60. onehot = paddle.transpose(onehot, (0, 3, 1, 2))
  61. return onehot
  62. def calculate_weights(self, label):
  63. hist = paddle.sum(label, axis=(1, 2)) * 1.0 / label.sum()
  64. hist = ((hist != 0) * self.upper_bound * (1 - hist)) + 1
  65. def custom_nll(self,
  66. logit,
  67. label,
  68. class_weights=None,
  69. border_weights=None,
  70. ignore_mask=None):
  71. soft = F.softmax(logit, axis=1)
  72. # calculate the valid soft where label is 1.
  73. soft_label = ((soft * label[:, :-1, :, :]).sum(
  74. 1, keepdim=True)) * (label[:, :-1, :, :].astype('float32'))
  75. soft = soft * (1 - label[:, :-1, :, :]) + soft_label
  76. logsoft = paddle.log(soft)
  77. if class_weights is not None:
  78. logsoft = class_weights.unsqueeze((0, 2, 3))
  79. logsoft = label[:, :-1, :, :] * logsoft
  80. logsoft = logsoft.sum(1)
  81. # border loss is divided equally
  82. logsoft = -1 / border_weights * logsoft * (1. - ignore_mask)
  83. n, _, h, w = label.shape
  84. logsoft = logsoft.sum() / (n * h * w - ignore_mask.sum() + 1)
  85. return logsoft
  86. def forward(self, logit, label):
  87. """
  88. Forward computation.
  89. Args:
  90. logit (Tensor): Logit tensor, the data type is float32, float64. Shape is
  91. (N, C), where C is number of classes, and if shape is more than 2D, this
  92. is (N, C, D1, D2,..., Dk), k >= 1.
  93. label (Tensor): Label tensor, the data type is int64. Shape is (N), where each
  94. value is 0 <= label[i] <= C-1, and if shape is more than 2D, this is
  95. (N, D1, D2,..., Dk), k >= 1.
  96. """
  97. n, c, h, w = logit.shape
  98. label.stop_gradient = True
  99. label = self.relax_onehot(label, c)
  100. weights = label[:, :-1, :, :].sum(1).astype('float32')
  101. ignore_mask = (weights == 0).astype('float32')
  102. # border is greater than 1, other is 1
  103. border_weights = weights + ignore_mask
  104. loss = 0
  105. class_weights = None
  106. for i in range(n):
  107. if self.calculate_weights:
  108. class_weights = self.calculate_weights(label[i])
  109. loss = loss + self.custom_nll(
  110. logit[i].unsqueeze(0),
  111. label[i].unsqueeze(0),
  112. class_weights=class_weights,
  113. border_weights=border_weights,
  114. ignore_mask=ignore_mask[i])
  115. return loss