| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141 |
- # 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
- import paddle.nn as nn
- import paddle.nn.functional as F
- from paddlex.paddleseg.cvlibs import manager
- @manager.LOSSES.add_component
- class DualTaskLoss(nn.Layer):
- """
- The dual task loss implement of GSCNN
- Args:
- ignore_index (int64): Specifies a target value that is ignored
- and does not contribute to the input gradient. Default ``255``.
- tau (float): the tau of gumbel softmax sample.
- """
- def __init__(self, ignore_index=255, tau=0.5):
- super().__init__()
- self.ignore_index = ignore_index
- self.tau = tau
- def _gumbel_softmax_sample(self, logit, tau=1, eps=1e-10):
- """
- Draw a sample from the Gumbel-Softmax distribution
- based on
- https://github.com/ericjang/gumbel-softmax/blob/3c8584924603869e90ca74ac20a6a03d99a91ef9/Categorical%20VAE.ipynb
- (MIT license)
- """
- gumbel_noise = paddle.rand(logit.shape)
- gumbel_noise = -paddle.log(eps - paddle.log(gumbel_noise + eps))
- logit = logit + gumbel_noise
- return F.softmax(logit / tau, axis=1)
- def compute_grad_mag(self, x):
- eps = 1e-6
- n, c, h, w = x.shape
- if h <= 1 or w <= 1:
- raise ValueError(
- 'The width and height of tensor to compute grad must be greater than 1, but the shape is {}.'
- .format(x.shape))
- x = self.conv_tri(x, r=4)
- kernel = [[-1, 0, 1]]
- kernel = paddle.to_tensor(kernel).astype('float32')
- kernel = 0.5 * kernel
- kernel_x = paddle.concat([kernel.unsqueeze((0, 1))] * c, axis=0)
- grad_x = F.conv2d(x, kernel_x, padding='same', groups=c)
- kernel_y = paddle.concat([kernel.t().unsqueeze((0, 1))] * c, axis=0)
- grad_y = F.conv2d(x, kernel_y, padding='same', groups=c)
- mag = paddle.sqrt(grad_x * grad_x + grad_y * grad_y + eps)
- return mag / mag.max()
- def conv_tri(self, input, r):
- """
- Convolves an image by a 2D triangle filter (the 1D triangle filter f is
- [1:r r+1 r:-1:1]/(r+1)^2, the 2D version is simply conv2(f,f'))
- """
- if r <= 1:
- raise ValueError(
- '`r` should be greater than 1, but it is {}.'.format(r))
- kernel = [
- list(range(1, r + 1)) + [r + 1] + list(reversed(range(1, r + 1)))
- ]
- kernel = paddle.to_tensor(kernel).astype('float32')
- kernel = kernel / (r + 1)**2
- input_ = F.pad(input, [1, 1, 0, 0], mode='replicate')
- input_ = F.pad(input_, [r, r, 0, 0], mode='reflect')
- input_ = [input_[:, :, :, :r], input, input_[:, :, :, -r:]]
- input_ = paddle.concat(input_, axis=3)
- tem = input_.clone()
- input_ = F.pad(input_, [0, 0, 1, 1], mode='replicate')
- input_ = F.pad(input_, [0, 0, r, r], mode='reflect')
- input_ = [input_[:, :, :r, :], tem, input_[:, :, -r:, :]]
- input_ = paddle.concat(input_, axis=2)
- c = input.shape[1]
- kernel_x = paddle.concat([kernel.unsqueeze((0, 1))] * c, axis=0)
- output = F.conv2d(input_, kernel_x, padding=0, groups=c)
- kernel_y = paddle.concat([kernel.t().unsqueeze((0, 1))] * c, axis=0)
- output = F.conv2d(output, kernel_y, padding=0, groups=c)
- return output
- def forward(self, logit, labels):
- # import pdb; pdb.set_trace()
- n, c, h, w = logit.shape
- th = 1e-8
- eps = 1e-10
- if len(labels.shape) == 3:
- labels = labels.unsqueeze(1)
- mask = (labels != self.ignore_index)
- mask.stop_gradient = True
- logit = logit * mask
- labels = labels * mask
- if len(labels.shape) == 4:
- labels = labels.squeeze(1)
- labels.stop_gradient = True
- labels = F.one_hot(labels, logit.shape[1]).transpose((0, 3, 1, 2))
- labels.stop_gradient = True
- g = self._gumbel_softmax_sample(logit, tau=self.tau)
- g = self.compute_grad_mag(g)
- g_hat = self.compute_grad_mag(labels)
- loss = F.l1_loss(g, g_hat, reduction='none')
- loss = loss * mask
- g_mask = (g > th).astype('float32')
- g_mask.stop_gradient = True
- g_mask_sum = paddle.sum(g_mask)
- loss_g = paddle.sum(loss * g_mask)
- if g_mask_sum > eps:
- loss_g = loss_g / g_mask_sum
- g_hat_mask = (g_hat > th).astype('float32')
- g_hat_mask.stop_gradient = True
- g_hat_mask_sum = paddle.sum(g_hat_mask)
- loss_g_hat = paddle.sum(loss * g_hat_mask)
- if g_hat_mask_sum > eps:
- loss_g_hat = loss_g_hat / g_hat_mask_sum
- total_loss = 0.5 * loss_g + 0.5 * loss_g_hat
- return total_loss
|