paddlex 2.0.0rc4__py3-none-any.whl → 3.0.0b2__py3-none-any.whl

paddlex/paddleseg/models/layers/nonlocal2d.py

@@ -1,155 +0,0 @@
-# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# 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.models import layers
-
-
-class NonLocal2D(nn.Layer):
-    """Basic Non-local module.
-    This model is the implementation of "Non-local Neural Networks"
-    (https://arxiv.org/abs/1711.07971)
-
-    Args:
-        in_channels (int): Channels of the input feature map.
-        reduction (int): Channel reduction ratio. Default: 2.
-        use_scale (bool): Whether to scale pairwise_weight by `1/sqrt(inter_channels)` when the mode is `embedded_gaussian`. Default: True.
-        sub_sample (bool): Whether to utilize max pooling after pairwise function. Default: False.
-        mode (str): Options are `gaussian`, `concatenation`, `embedded_gaussian` and `dot_product`. Default: embedded_gaussian.
-    """
-
-    def __init__(self,
-                 in_channels,
-                 reduction=2,
-                 use_scale=True,
-                 sub_sample=False,
-                 mode='embedded_gaussian'):
-        super(NonLocal2D, self).__init__()
-        self.in_channels = in_channels
-        self.reduction = reduction
-        self.use_scale = use_scale
-        self.sub_sample = sub_sample
-        self.mode = mode
-        if mode not in [
-                'gaussian', 'embedded_gaussian', 'dot_product', 'concatenation'
-        ]:
-            raise ValueError(
-                "Mode should be in 'gaussian', 'concatenation','embedded_gaussian' or 'dot_product'."
-            )
-
-        self.inter_channels = max(in_channels // reduction, 1)
-
-        self.g = nn.Conv2D(
-            in_channels=self.in_channels,
-            out_channels=self.inter_channels,
-            kernel_size=1)
-        self.conv_out = layers.ConvBNReLU(
-            in_channels=self.inter_channels,
-            out_channels=self.in_channels,
-            kernel_size=1,
-            bias_attr=False)
-
-        if self.mode != "gaussian":
-            self.theta = nn.Conv2D(
-                in_channels=self.in_channels,
-                out_channels=self.inter_channels,
-                kernel_size=1)
-            self.phi = nn.Conv2D(
-                in_channels=self.in_channels,
-                out_channels=self.inter_channels,
-                kernel_size=1)
-
-        if self.mode == "concatenation":
-            self.concat_project = layers.ConvBNReLU(
-                in_channels=self.inter_channels * 2,
-                out_channels=1,
-                kernel_size=1,
-                bias_attr=False)
-
-        if self.sub_sample:
-            max_pool_layer = nn.MaxPool2D(kernel_size=(2, 2))
-            self.g = nn.Sequential(self.g, max_pool_layer)
-            if self.mode != 'gaussian':
-                self.phi = nn.Sequential(self.phi, max_pool_layer)
-            else:
-                self.phi = max_pool_layer
-
-    def gaussian(self, theta_x, phi_x):
-        pairwise_weight = paddle.matmul(theta_x, phi_x)
-        pairwise_weight = F.softmax(pairwise_weight, axis=-1)
-        return pairwise_weight
-
-    def embedded_gaussian(self, theta_x, phi_x):
-        pairwise_weight = paddle.matmul(theta_x, phi_x)
-        if self.use_scale:
-            pairwise_weight /= theta_x.shape[-1]**0.5
-        pairwise_weight = F.softmax(pairwise_weight, -1)
-        return pairwise_weight
-
-    def dot_product(self, theta_x, phi_x):
-        pairwise_weight = paddle.matmul(theta_x, phi_x)
-        pairwise_weight /= pairwise_weight.shape[-1]
-        return pairwise_weight
-
-    def concatenation(self, theta_x, phi_x):
-        h = theta_x.shape[2]
-        w = phi_x.shape[3]
-        theta_x = paddle.tile(theta_x, [1, 1, 1, w])
-        phi_x = paddle.tile(phi_x, [1, 1, h, 1])
-
-        concat_feature = paddle.concat([theta_x, phi_x], axis=1)
-        pairwise_weight = self.concat_project(concat_feature)
-        n, _, h, w = pairwise_weight.shape
-        pairwise_weight = paddle.reshape(pairwise_weight, [n, h, w])
-        pairwise_weight /= pairwise_weight.shape[-1]
-        return pairwise_weight
-
-    def forward(self, x):
-        n, c, h, w = x.shape
-        g_x = paddle.reshape(self.g(x), [n, self.inter_channels, -1])
-        g_x = paddle.transpose(g_x, [0, 2, 1])
-
-        if self.mode == 'gaussian':
-            theta_x = paddle.reshape(x, [n, self.inter_channels, -1])
-            theta_x = paddle.transpose(theta_x, [0, 2, 1])
-            if self.sub_sample:
-                phi_x = paddle.reshape(
-                    self.phi(x), [n, self.inter_channels, -1])
-            else:
-                phi_x = paddle.reshape(x, [n, self.in_channels, -1])
-
-        elif self.mode == 'concatenation':
-            theta_x = paddle.reshape(
-                self.theta(x), [n, self.inter_channels, -1, 1])
-            phi_x = paddle.reshape(
-                self.phi(x), [n, self.inter_channels, 1, -1])
-
-        else:
-            theta_x = paddle.reshape(
-                self.theta(x), [n, self.inter_channels, -1])
-            theta_x = paddle.transpose(theta_x, [0, 2, 1])
-            phi_x = paddle.reshape(self.phi(x), [n, self.inter_channels, -1])
-
-        pairwise_func = getattr(self, self.mode)
-        pairwise_weight = pairwise_func(theta_x, phi_x)
-        y = paddle.matmul(pairwise_weight, g_x)
-        y = paddle.transpose(y, [0, 2, 1])
-        y = paddle.reshape(y, [n, self.inter_channels, h, w])
-
-        output = x + self.conv_out(y)
-
-        return output

paddlex/paddleseg/models/layers/pyramid_pool.py

@@ -1,182 +0,0 @@
-# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# 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.functional as F
-from paddle import nn
-
-from paddlex.paddleseg.models import layers
-
-
-class ASPPModule(nn.Layer):
-    """
-    Atrous Spatial Pyramid Pooling.
-
-    Args:
-        aspp_ratios (tuple): The dilation rate using in ASSP module.
-        in_channels (int): The number of input channels.
-        out_channels (int): The number of output channels.
-        align_corners (bool): An argument of F.interpolate. It should be set to False when the output size of feature
-            is even, e.g. 1024x512, otherwise it is True, e.g. 769x769.
-        use_sep_conv (bool, optional): If using separable conv in ASPP module. Default: False.
-        image_pooling (bool, optional): If augmented with image-level features. Default: False
-    """
-
-    def __init__(self,
-                 aspp_ratios,
-                 in_channels,
-                 out_channels,
-                 align_corners,
-                 use_sep_conv=False,
-                 image_pooling=False):
-        super().__init__()
-
-        self.align_corners = align_corners
-        self.aspp_blocks = nn.LayerList()
-
-        for ratio in aspp_ratios:
-            if use_sep_conv and ratio > 1:
-                conv_func = layers.SeparableConvBNReLU
-            else:
-                conv_func = layers.ConvBNReLU
-
-            block = conv_func(
-                in_channels=in_channels,
-                out_channels=out_channels,
-                kernel_size=1 if ratio == 1 else 3,
-                dilation=ratio,
-                padding=0 if ratio == 1 else ratio)
-            self.aspp_blocks.append(block)
-
-        out_size = len(self.aspp_blocks)
-
-        if image_pooling:
-            self.global_avg_pool = nn.Sequential(
-                nn.AdaptiveAvgPool2D(output_size=(1, 1)),
-                layers.ConvBNReLU(
-                    in_channels, out_channels, kernel_size=1, bias_attr=False))
-            out_size += 1
-        self.image_pooling = image_pooling
-
-        self.conv_bn_relu = layers.ConvBNReLU(
-            in_channels=out_channels * out_size,
-            out_channels=out_channels,
-            kernel_size=1)
-
-        self.dropout = nn.Dropout(p=0.1)  # drop rate
-
-    def forward(self, x):
-        outputs = []
-        interpolate_shape = paddle.shape(x)[2:]
-        for block in self.aspp_blocks:
-            y = block(x)
-            y = F.interpolate(
-                y,
-                interpolate_shape,
-                mode='bilinear',
-                align_corners=self.align_corners)
-            outputs.append(y)
-
-        if self.image_pooling:
-            img_avg = self.global_avg_pool(x)
-            img_avg = F.interpolate(
-                img_avg,
-                interpolate_shape,
-                mode='bilinear',
-                align_corners=self.align_corners)
-            outputs.append(img_avg)
-
-        x = paddle.concat(outputs, axis=1)
-        x = self.conv_bn_relu(x)
-        x = self.dropout(x)
-
-        return x
-
-
-class PPModule(nn.Layer):
-    """
-    Pyramid pooling module originally in PSPNet.
-
-    Args:
-        in_channels (int): The number of intput channels to pyramid pooling module.
-        out_channels (int): The number of output channels after pyramid pooling module.
-        bin_sizes (tuple, optional): The out size of pooled feature maps. Default: (1, 2, 3, 6).
-        dim_reduction (bool, optional): A bool value represents if reducing dimension after pooling. Default: True.
-        align_corners (bool): An argument of F.interpolate. It should be set to False when the output size of feature
-            is even, e.g. 1024x512, otherwise it is True, e.g. 769x769.
-    """
-
-    def __init__(self, in_channels, out_channels, bin_sizes, dim_reduction,
-                 align_corners):
-        super().__init__()
-
-        self.bin_sizes = bin_sizes
-
-        inter_channels = in_channels
-        if dim_reduction:
-            inter_channels = in_channels // len(bin_sizes)
-
-        # we use dimension reduction after pooling mentioned in original implementation.
-        self.stages = nn.LayerList([
-            self._make_stage(in_channels, inter_channels, size)
-            for size in bin_sizes
-        ])
-
-        self.conv_bn_relu2 = layers.ConvBNReLU(
-            in_channels=in_channels + inter_channels * len(bin_sizes),
-            out_channels=out_channels,
-            kernel_size=3,
-            padding=1)
-
-        self.align_corners = align_corners
-
-    def _make_stage(self, in_channels, out_channels, size):
-        """
-        Create one pooling layer.
-
-        In our implementation, we adopt the same dimension reduction as the original paper that might be
-        slightly different with other implementations.
-
-        After pooling, the channels are reduced to 1/len(bin_sizes) immediately, while some other implementations
-        keep the channels to be same.
-
-        Args:
-            in_channels (int): The number of intput channels to pyramid pooling module.
-            size (int): The out size of the pooled layer.
-
-        Returns:
-            conv (Tensor): A tensor after Pyramid Pooling Module.
-        """
-
-        prior = nn.AdaptiveAvgPool2D(output_size=(size, size))
-        conv = layers.ConvBNReLU(
-            in_channels=in_channels, out_channels=out_channels, kernel_size=1)
-
-        return nn.Sequential(prior, conv)
-
-    def forward(self, input):
-        cat_layers = []
-        for stage in self.stages:
-            x = stage(input)
-            x = F.interpolate(
-                x,
-                paddle.shape(input)[2:],
-                mode='bilinear',
-                align_corners=self.align_corners)
-            cat_layers.append(x)
-        cat_layers = [input] + cat_layers[::-1]
-        cat = paddle.concat(cat_layers, axis=1)
-        out = self.conv_bn_relu2(cat)
-
-        return out

paddlex/paddleseg/models/losses/__init__.py

@@ -1,27 +0,0 @@
-# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# 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.
-
-from .mixed_loss import MixedLoss
-from .cross_entropy_loss import CrossEntropyLoss
-from .binary_cross_entropy_loss import BCELoss
-from .lovasz_loss import LovaszSoftmaxLoss, LovaszHingeLoss
-from .gscnn_dual_task_loss import DualTaskLoss
-from .edge_attention_loss import EdgeAttentionLoss
-from .bootstrapped_cross_entropy import BootstrappedCrossEntropyLoss
-from .dice_loss import DiceLoss
-from .ohem_cross_entropy_loss import OhemCrossEntropyLoss
-from .decoupledsegnet_relax_boundary_loss import RelaxBoundaryLoss
-from .ohem_edge_attention_loss import OhemEdgeAttentionLoss
-from .l1_loss import L1Loss
-from .mean_square_error_loss import MSELoss

paddlex/paddleseg/models/losses/binary_cross_entropy_loss.py

@@ -1,174 +0,0 @@
-# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# 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 BCELoss(nn.Layer):
-    """
-    This operator combines the sigmoid layer and the :ref:`api_nn_loss_BCELoss` layer.
-    Also, we can see it as the combine of ``sigmoid_cross_entropy_with_logits``
-    layer and some reduce operations.
-    This measures the element-wise probability error in classification tasks
-    in which each class is independent.
-    This can be thought of as predicting labels for a data-point, where labels
-    are not mutually exclusive. For example, a news article can be about
-    politics, technology or sports at the same time or none of these.
-    First this operator calculate loss function as follows:
-    .. math::
-           Out = -Labels * \\log(\\sigma(Logit)) - (1 - Labels) * \\log(1 - \\sigma(Logit))
-    We know that :math:`\\sigma(Logit) = \\frac{1}{1 + \\e^{-Logit}}`. By substituting this we get:
-    .. math::
-           Out = Logit - Logit * Labels + \\log(1 + \\e^{-Logit})
-    For stability and to prevent overflow of :math:`\\e^{-Logit}` when Logit < 0,
-    we reformulate the loss as follows:
-    .. math::
-           Out = \\max(Logit, 0) - Logit * Labels + \\log(1 + \\e^{-\|Logit\|})
-    Then, if ``weight`` or ``pos_weight`` is not None, this operator multiply the
-    weight tensor on the loss `Out`. The ``weight`` tensor will attach different
-    weight on every items in the batch. The ``pos_weight`` will attach different
-    weight on the positive label of each class.
-    Finally, this operator applies reduce operation on the loss.
-    If :attr:`reduction` set to ``'none'``, the operator will return the original loss `Out`.
-    If :attr:`reduction` set to ``'mean'``, the reduced mean loss is :math:`Out = MEAN(Out)`.
-    If :attr:`reduction` set to ``'sum'``, the reduced sum loss is :math:`Out = SUM(Out)`.
-    Note that the target labels ``label`` should be numbers between 0 and 1.
-    Args:
-        weight (Tensor | str, optional): A manual rescaling weight given to the loss of each
-            batch element. If given, it has to be a 1D Tensor whose size is `[N, ]`,
-            The data type is float32, float64. If type is str, it should equal to 'dynamic'.
-            It will compute weight dynamically in every step.
-            Default is ``'None'``.
-        pos_weight (float|str, optional): A weight of positive examples. If type is str,
-            it should equal to 'dynamic'. It will compute weight dynamically in every step.
-            Default is ``'None'``.
-        ignore_index (int64, optional): Specifies a target value that is ignored
-            and does not contribute to the input gradient. Default ``255``.
-        edge_label (bool, optional): Whether to use edge label. Default: False
-    Shapes:
-        logit (Tensor): The input predications tensor. 2-D tensor with shape: [N, *],
-            N is batch_size, `*` means number of additional dimensions. The ``logit``
-            is usually the output of Linear layer. Available dtype is float32, float64.
-        label (Tensor): The target labels tensor. 2-D tensor with the same shape as
-            ``logit``. The target labels which values should be numbers between 0 and 1.
-            Available dtype is float32, float64.
-    Returns:
-        A callable object of BCEWithLogitsLoss.
-    Examples:
-        .. code-block:: python
-            import paddle
-            paddle.disable_static()
-            logit = paddle.to_tensor([5.0, 1.0, 3.0], dtype="float32")
-            label = paddle.to_tensor([1.0, 0.0, 1.0], dtype="float32")
-            bce_logit_loss = paddle.nn.BCEWithLogitsLoss()
-            output = bce_logit_loss(logit, label)
-            print(output.numpy())  # [0.45618808]
-    """
-
-    def __init__(self,
-                 weight=None,
-                 pos_weight=None,
-                 ignore_index=255,
-                 edge_label=False):
-        super().__init__()
-        self.weight = weight
-        self.pos_weight = pos_weight
-        self.ignore_index = ignore_index
-        self.edge_label = edge_label
-        self.EPS = 1e-10
-
-        if self.weight is not None:
-            if isinstance(self.weight, str):
-                if self.weight != 'dynamic':
-                    raise ValueError(
-                        "if type of `weight` is str, it should equal to 'dynamic', but it is {}"
-                        .format(self.weight))
-            elif isinstance(self.weight, paddle.VarBase):
-                raise TypeError(
-                    'The type of `weight` is wrong, it should be Tensor or str, but it is {}'
-                    .format(type(self.weight)))
-
-        if self.pos_weight is not None:
-            if isinstance(self.pos_weight, str):
-                if self.pos_weight != 'dynamic':
-                    raise ValueError(
-                        "if type of `pos_weight` is str, it should equal to 'dynamic', but it is {}"
-                        .format(self.pos_weight))
-            elif isinstance(self.pos_weight, float):
-                self.pos_weight = paddle.to_tensor(
-                    self.pos_weight, dtype='float32')
-            else:
-                raise TypeError(
-                    'The type of `pos_weight` is wrong, it should be float or str, but it is {}'
-                    .format(type(self.pos_weight)))
-
-    def forward(self, logit, label):
-        """
-        Forward computation.
-
-        Args:
-            logit (Tensor): Logit tensor, the data type is float32, float64. Shape is
-                (N, C), where C is number of classes, and if shape is more than 2D, this
-                is (N, C, D1, D2,..., Dk), k >= 1.
-            label (Tensor): Label tensor, the data type is int64. Shape is (N, C), where each
-                value is 0 or 1, and if shape is more than 2D, this is
-                (N, C, D1, D2,..., Dk), k >= 1.
-        """
-        if len(label.shape) != len(logit.shape):
-            label = paddle.unsqueeze(label, 1)
-        mask = (label != self.ignore_index)
-        mask = paddle.cast(mask, 'float32')
-        # label.shape should equal to the logit.shape
-        if label.shape[1] != logit.shape[1]:
-            label = label.squeeze(1)
-            label = F.one_hot(label, logit.shape[1])
-            label = label.transpose((0, 3, 1, 2))
-        if isinstance(self.weight, str):
-            pos_index = (label == 1)
-            neg_index = (label == 0)
-            pos_num = paddle.sum(pos_index.astype('float32'))
-            neg_num = paddle.sum(neg_index.astype('float32'))
-            sum_num = pos_num + neg_num
-            weight_pos = 2 * neg_num / (sum_num + self.EPS)
-            weight_neg = 2 * pos_num / (sum_num + self.EPS)
-            weight = weight_pos * label + weight_neg * (1 - label)
-        else:
-            weight = self.weight
-        if isinstance(self.pos_weight, str):
-            pos_index = (label == 1)
-            neg_index = (label == 0)
-            pos_num = paddle.sum(pos_index.astype('float32'))
-            neg_num = paddle.sum(neg_index.astype('float32'))
-            sum_num = pos_num + neg_num
-            pos_weight = 2 * neg_num / (sum_num + self.EPS)
-        else:
-            pos_weight = self.pos_weight
-        label = label.astype('float32')
-        loss = paddle.nn.functional.binary_cross_entropy_with_logits(
-            logit,
-            label,
-            weight=weight,
-            reduction='none',
-            pos_weight=pos_weight)
-        loss = loss * mask
-        loss = paddle.mean(loss) / (paddle.mean(mask) + self.EPS)
-        label.stop_gradient = True
-        mask.stop_gradient = True
-
-        return loss

paddlex/paddleseg/models/losses/bootstrapped_cross_entropy.py

@@ -1,73 +0,0 @@
-# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# 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 paddlex.paddleseg.cvlibs import manager
-
-
-@manager.LOSSES.add_component
-class BootstrappedCrossEntropyLoss(nn.Layer):
-    """
-    Implements the cross entropy loss function.
-
-    Args:
-        min_K (int): the minimum number of pixels to be counted in loss computation.
-        loss_th (float): the loss threshold. Only loss that is larger than the threshold
-            would be calculated.
-        weight (tuple|list, optional): The weight for different classes. Default: None.
-        ignore_index (int, optional): Specifies a target value that is ignored
-            and does not contribute to the input gradient. Default: 255.
-    """
-
-    def __init__(self, min_K, loss_th, weight=None, ignore_index=255):
-        super().__init__()
-        self.ignore_index = ignore_index
-        self.K = min_K
-        self.threshold = loss_th
-        if weight is not None:
-            weight = paddle.to_tensor(weight, dtype='float32')
-        self.weight = weight
-
-    def forward(self, logit, label):
-
-        n, c, h, w = logit.shape
-        total_loss = 0.0
-        if len(label.shape) != len(logit.shape):
-            label = paddle.unsqueeze(label, 1)
-
-        for i in range(n):
-            x = paddle.unsqueeze(logit[i], 0)
-            y = paddle.unsqueeze(label[i], 0)
-            x = paddle.transpose(x, (0, 2, 3, 1))
-            y = paddle.transpose(y, (0, 2, 3, 1))
-            x = paddle.reshape(x, shape=(-1, c))
-            y = paddle.reshape(y, shape=(-1, ))
-            loss = F.cross_entropy(
-                x,
-                y,
-                weight=self.weight,
-                ignore_index=self.ignore_index,
-                reduction="none")
-            sorted_loss = paddle.sort(loss, descending=True)
-            if sorted_loss[self.K] > self.threshold:
-                new_indices = paddle.nonzero(sorted_loss > self.threshold)
-                loss = paddle.gather(sorted_loss, new_indices)
-            else:
-                loss = sorted_loss[:self.K]
-
-            total_loss += paddle.mean(loss)
-        return total_loss / float(n)