jde_loss.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182
  1. # Copyright (c) 2021 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. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import paddle
  18. import paddle.nn as nn
  19. import paddle.nn.functional as F
  20. from paddlex.ppdet.core.workspace import register
  21. __all__ = ['JDEDetectionLoss', 'JDEEmbeddingLoss', 'JDELoss']
  22. @register
  23. class JDEDetectionLoss(nn.Layer):
  24. __shared__ = ['num_classes']
  25. def __init__(self, num_classes=1):
  26. super(JDEDetectionLoss, self).__init__()
  27. self.num_classes = num_classes
  28. def det_loss(self, p_det, anchor, t_conf, t_box):
  29. pshape = paddle.shape(p_det)
  30. pshape.stop_gradient = True
  31. nB, nGh, nGw = pshape[0], pshape[-2], pshape[-1]
  32. nA = len(anchor)
  33. p_det = paddle.reshape(
  34. p_det, [nB, nA, self.num_classes + 5, nGh, nGw]).transpose(
  35. (0, 1, 3, 4, 2))
  36. # 1. loss_conf: cross_entropy
  37. p_conf = p_det[:, :, :, :, 4:6]
  38. p_conf_flatten = paddle.reshape(p_conf, [-1, 2])
  39. t_conf_flatten = t_conf.flatten()
  40. t_conf_flatten = paddle.cast(t_conf_flatten, dtype="int64")
  41. t_conf_flatten.stop_gradient = True
  42. loss_conf = F.cross_entropy(
  43. p_conf_flatten, t_conf_flatten, ignore_index=-1, reduction='mean')
  44. loss_conf.stop_gradient = False
  45. # 2. loss_box: smooth_l1_loss
  46. p_box = p_det[:, :, :, :, :4]
  47. p_box_flatten = paddle.reshape(p_box, [-1, 4])
  48. t_box_flatten = paddle.reshape(t_box, [-1, 4])
  49. fg_inds = paddle.nonzero(t_conf_flatten > 0).flatten()
  50. if fg_inds.numel() > 0:
  51. reg_delta = paddle.gather(p_box_flatten, fg_inds)
  52. reg_target = paddle.gather(t_box_flatten, fg_inds)
  53. else:
  54. reg_delta = paddle.to_tensor([0, 0, 0, 0], dtype='float32')
  55. reg_delta.stop_gradient = False
  56. reg_target = paddle.to_tensor([0, 0, 0, 0], dtype='float32')
  57. reg_target.stop_gradient = True
  58. loss_box = F.smooth_l1_loss(
  59. reg_delta, reg_target, reduction='mean', delta=1.0)
  60. loss_box.stop_gradient = False
  61. return loss_conf, loss_box
  62. def forward(self, det_outs, targets, anchors):
  63. """
  64. Args:
  65. det_outs (list[Tensor]): output from detection head, each one
  66. is a 4-D Tensor with shape [N, C, H, W].
  67. targets (dict): contains 'im_id', 'gt_bbox', 'gt_ide', 'image',
  68. 'im_shape', 'scale_factor' and 'tbox', 'tconf', 'tide' of
  69. each FPN level.
  70. anchors (list[list]): anchor setting of JDE model, N row M col, N is
  71. the anchor levels(FPN levels), M is the anchor scales each
  72. level.
  73. """
  74. assert len(det_outs) == len(anchors)
  75. loss_confs = []
  76. loss_boxes = []
  77. for i, (p_det, anchor) in enumerate(zip(det_outs, anchors)):
  78. t_conf = targets['tconf{}'.format(i)]
  79. t_box = targets['tbox{}'.format(i)]
  80. loss_conf, loss_box = self.det_loss(p_det, anchor, t_conf, t_box)
  81. loss_confs.append(loss_conf)
  82. loss_boxes.append(loss_box)
  83. return {'loss_confs': loss_confs, 'loss_boxes': loss_boxes}
  84. @register
  85. class JDEEmbeddingLoss(nn.Layer):
  86. def __init__(self, ):
  87. super(JDEEmbeddingLoss, self).__init__()
  88. self.phony = self.create_parameter(shape=[1], dtype="float32")
  89. def emb_loss(self, p_ide, t_conf, t_ide, emb_scale, classifier):
  90. emb_dim = p_ide.shape[1]
  91. p_ide = p_ide.transpose((0, 2, 3, 1))
  92. p_ide_flatten = paddle.reshape(p_ide, [-1, emb_dim])
  93. mask = t_conf > 0
  94. mask = paddle.cast(mask, dtype="int64")
  95. mask.stop_gradient = True
  96. emb_mask = mask.max(1).flatten()
  97. emb_mask_inds = paddle.nonzero(emb_mask > 0).flatten()
  98. emb_mask_inds.stop_gradient = True
  99. # use max(1) to decide the id, TODO: more reseanable strategy
  100. t_ide_flatten = t_ide.max(1).flatten()
  101. t_ide_flatten = paddle.cast(t_ide_flatten, dtype="int64")
  102. valid_inds = paddle.nonzero(t_ide_flatten != -1).flatten()
  103. if emb_mask_inds.numel() == 0 or valid_inds.numel() == 0:
  104. # loss_ide = paddle.to_tensor([0]) # will be error in gradient backward
  105. loss_ide = self.phony * 0 # todo
  106. else:
  107. embedding = paddle.gather(p_ide_flatten, emb_mask_inds)
  108. embedding = emb_scale * F.normalize(embedding)
  109. logits = classifier(embedding)
  110. ide_target = paddle.gather(t_ide_flatten, emb_mask_inds)
  111. loss_ide = F.cross_entropy(
  112. logits, ide_target, ignore_index=-1, reduction='mean')
  113. loss_ide.stop_gradient = False
  114. return loss_ide
  115. def forward(self, ide_outs, targets, emb_scale, classifier):
  116. loss_ides = []
  117. for i, p_ide in enumerate(ide_outs):
  118. t_conf = targets['tconf{}'.format(i)]
  119. t_ide = targets['tide{}'.format(i)]
  120. loss_ide = self.emb_loss(p_ide, t_conf, t_ide, emb_scale,
  121. classifier)
  122. loss_ides.append(loss_ide)
  123. return loss_ides
  124. @register
  125. class JDELoss(nn.Layer):
  126. def __init__(self):
  127. super(JDELoss, self).__init__()
  128. def forward(self, loss_confs, loss_boxes, loss_ides, loss_params_cls,
  129. loss_params_reg, loss_params_ide, targets):
  130. assert len(loss_confs) == len(loss_boxes) == len(loss_ides)
  131. assert len(loss_params_cls) == len(loss_params_reg) == len(
  132. loss_params_ide)
  133. assert len(loss_confs) == len(loss_params_cls)
  134. batchsize = targets['gt_bbox'].shape[0]
  135. nTargets = paddle.nonzero(paddle.sum(targets['gt_bbox'],
  136. axis=2)).shape[0] / batchsize
  137. nTargets = paddle.to_tensor(nTargets, dtype='float32')
  138. nTargets.stop_gradient = True
  139. jde_losses = []
  140. for i, (loss_conf, loss_box, loss_ide, l_conf_p, l_box_p,
  141. l_ide_p) in enumerate(
  142. zip(loss_confs, loss_boxes, loss_ides, loss_params_cls,
  143. loss_params_reg, loss_params_ide)):
  144. jde_loss = l_conf_p(loss_conf) + l_box_p(loss_box) + l_ide_p(
  145. loss_ide)
  146. jde_losses.append(jde_loss)
  147. loss_all = {
  148. "loss_conf": sum(loss_confs),
  149. "loss_box": sum(loss_boxes),
  150. "loss_ide": sum(loss_ides),
  151. "loss": sum(jde_losses),
  152. "nTargets": nTargets,
  153. }
  154. return loss_all