• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Huawei Technologies Co., Ltd
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# ============================================================================
15"""YOLOV3 loss."""
16from mindspore.ops import operations as P
17import mindspore.nn as nn
18
19
20class XYLoss(nn.Cell):
21    """Loss for x and y."""
22    def __init__(self):
23        super(XYLoss, self).__init__()
24        self.cross_entropy = P.SigmoidCrossEntropyWithLogits()
25        self.reduce_sum = P.ReduceSum()
26
27    def construct(self, object_mask, box_loss_scale, predict_xy, true_xy):
28        xy_loss = object_mask * box_loss_scale * self.cross_entropy(predict_xy, true_xy)
29        xy_loss = self.reduce_sum(xy_loss, ())
30        return xy_loss
31
32
33class WHLoss(nn.Cell):
34    """Loss for w and h."""
35    def __init__(self):
36        super(WHLoss, self).__init__()
37        self.square = P.Square()
38        self.reduce_sum = P.ReduceSum()
39
40    def construct(self, object_mask, box_loss_scale, predict_wh, true_wh):
41        wh_loss = object_mask * box_loss_scale * 0.5 * P.Square()(true_wh - predict_wh)
42        wh_loss = self.reduce_sum(wh_loss, ())
43        return wh_loss
44
45
46class ConfidenceLoss(nn.Cell):
47    """Loss for confidence."""
48    def __init__(self):
49        super(ConfidenceLoss, self).__init__()
50        self.cross_entropy = P.SigmoidCrossEntropyWithLogits()
51        self.reduce_sum = P.ReduceSum()
52
53    def construct(self, object_mask, predict_confidence, ignore_mask):
54        confidence_loss = self.cross_entropy(predict_confidence, object_mask)
55        confidence_loss = object_mask * confidence_loss + (1 - object_mask) * confidence_loss * ignore_mask
56        confidence_loss = self.reduce_sum(confidence_loss, ())
57        return confidence_loss
58
59
60class ClassLoss(nn.Cell):
61    """Loss for classification."""
62    def __init__(self):
63        super(ClassLoss, self).__init__()
64        self.cross_entropy = P.SigmoidCrossEntropyWithLogits()
65        self.reduce_sum = P.ReduceSum()
66
67    def construct(self, object_mask, predict_class, class_probs):
68        class_loss = object_mask * self.cross_entropy(predict_class, class_probs)
69        class_loss = self.reduce_sum(class_loss, ())
70        return class_loss
71