# Copyright (c) 2023 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 __future__ import absolute_import from __future__ import division from __future__ import print_function import paddle import paddle.nn as nn import paddle.nn.functional as F from ppdet.core.workspace import register from ppdet.modeling.losses.iou_loss import GIoULoss from .utils import QFLv2 from ppdet.utils.logger import setup_logger logger = setup_logger(__name__) __all__ = [ 'SSODFCOSLoss', 'SSODPPYOLOELoss', ] @register class SSODFCOSLoss(nn.Layer): def __init__(self, loss_weight=1.0): super(SSODFCOSLoss, self).__init__() self.loss_weight = loss_weight def forward(self, student_head_outs, teacher_head_outs, train_cfg): # for semi-det distill student_logits, student_deltas, student_quality = student_head_outs teacher_logits, teacher_deltas, teacher_quality = teacher_head_outs nc = student_logits[0].shape[1] student_logits = paddle.concat( [ _.transpose([0, 2, 3, 1]).reshape([-1, nc]) for _ in student_logits ], axis=0) teacher_logits = paddle.concat( [ _.transpose([0, 2, 3, 1]).reshape([-1, nc]) for _ in teacher_logits ], axis=0) student_deltas = paddle.concat( [ _.transpose([0, 2, 3, 1]).reshape([-1, 4]) for _ in student_deltas ], axis=0) teacher_deltas = paddle.concat( [ _.transpose([0, 2, 3, 1]).reshape([-1, 4]) for _ in teacher_deltas ], axis=0) student_quality = paddle.concat( [ _.transpose([0, 2, 3, 1]).reshape([-1, 1]) for _ in student_quality ], axis=0) teacher_quality = paddle.concat( [ _.transpose([0, 2, 3, 1]).reshape([-1, 1]) for _ in teacher_quality ], axis=0) ratio = train_cfg.get('ratio', 0.01) with paddle.no_grad(): # Region Selection count_num = int(teacher_logits.shape[0] * ratio) teacher_probs = F.sigmoid(teacher_logits) max_vals = paddle.max(teacher_probs, 1) sorted_vals, sorted_inds = paddle.topk(max_vals, teacher_logits.shape[0]) mask = paddle.zeros_like(max_vals) mask[sorted_inds[:count_num]] = 1. fg_num = sorted_vals[:count_num].sum() b_mask = mask > 0 # distill_loss_cls loss_logits = QFLv2( F.sigmoid(student_logits), teacher_probs, weight=mask, reduction="sum") / fg_num # distill_loss_box inputs = paddle.concat( (-student_deltas[b_mask][..., :2], student_deltas[b_mask][..., 2:]), axis=-1) targets = paddle.concat( (-teacher_deltas[b_mask][..., :2], teacher_deltas[b_mask][..., 2:]), axis=-1) iou_loss = GIoULoss(reduction='mean') loss_deltas = iou_loss(inputs, targets) # distill_loss_quality loss_quality = F.binary_cross_entropy( F.sigmoid(student_quality[b_mask]), F.sigmoid(teacher_quality[b_mask]), reduction='mean') return { "distill_loss_cls": loss_logits, "distill_loss_box": loss_deltas, "distill_loss_quality": loss_quality, "fg_sum": fg_num, } @register class SSODPPYOLOELoss(nn.Layer): def __init__(self, loss_weight=1.0): super(SSODPPYOLOELoss, self).__init__() self.loss_weight = loss_weight def forward(self, student_head_outs, teacher_head_outs, train_cfg): # for semi-det distill # student_probs: already sigmoid student_probs, student_deltas, student_dfl = student_head_outs teacher_probs, teacher_deltas, teacher_dfl = teacher_head_outs bs, l, nc = student_probs.shape[:] # bs, l, num_classes bs, l, _, reg_ch = student_dfl.shape[:] # bs, l, 4, reg_ch student_probs = student_probs.reshape([-1, nc]) teacher_probs = teacher_probs.reshape([-1, nc]) student_deltas = student_deltas.reshape([-1, 4]) teacher_deltas = teacher_deltas.reshape([-1, 4]) student_dfl = student_dfl.reshape([-1, 4, reg_ch]) teacher_dfl = teacher_dfl.reshape([-1, 4, reg_ch]) ratio = train_cfg.get('ratio', 0.01) # for contrast loss curr_iter = train_cfg['curr_iter'] st_iter = train_cfg['st_iter'] if curr_iter == st_iter + 1: # start semi-det training self.queue_ptr = 0 self.queue_size = int(bs * l * ratio) self.queue_feats = paddle.zeros([self.queue_size, nc]) self.queue_probs = paddle.zeros([self.queue_size, nc]) contrast_loss_cfg = train_cfg['contrast_loss'] temperature = contrast_loss_cfg.get('temperature', 0.2) alpha = contrast_loss_cfg.get('alpha', 0.9) smooth_iter = contrast_loss_cfg.get('smooth_iter', 100) + st_iter with paddle.no_grad(): # Region Selection count_num = int(teacher_probs.shape[0] * ratio) max_vals = paddle.max(teacher_probs, 1) sorted_vals, sorted_inds = paddle.topk(max_vals, teacher_probs.shape[0]) mask = paddle.zeros_like(max_vals) mask[sorted_inds[:count_num]] = 1. fg_num = sorted_vals[:count_num].sum() b_mask = mask > 0. # for contrast loss probs = teacher_probs[b_mask].detach() if curr_iter > smooth_iter: # memory-smoothing A = paddle.exp( paddle.mm(teacher_probs[b_mask], self.queue_probs.t()) / temperature) A = A / A.sum(1, keepdim=True) probs = alpha * probs + (1 - alpha) * paddle.mm( A, self.queue_probs) n = student_probs[b_mask].shape[0] # update memory bank self.queue_feats[self.queue_ptr:self.queue_ptr + n, :] = teacher_probs[b_mask].detach() self.queue_probs[self.queue_ptr:self.queue_ptr + n, :] = teacher_probs[b_mask].detach() self.queue_ptr = (self.queue_ptr + n) % self.queue_size # embedding similarity sim = paddle.exp( paddle.mm(student_probs[b_mask], teacher_probs[b_mask].t()) / 0.2) sim_probs = sim / sim.sum(1, keepdim=True) # pseudo-label graph with self-loop Q = paddle.mm(probs, probs.t()) Q.fill_diagonal_(1) pos_mask = (Q >= 0.5).astype('float32') Q = Q * pos_mask Q = Q / Q.sum(1, keepdim=True) # contrastive loss loss_contrast = -(paddle.log(sim_probs + 1e-7) * Q).sum(1) loss_contrast = loss_contrast.mean() # distill_loss_cls loss_cls = QFLv2( student_probs, teacher_probs, weight=mask, reduction="sum") / fg_num # distill_loss_iou inputs = paddle.concat( (-student_deltas[b_mask][..., :2], student_deltas[b_mask][..., 2:]), -1) targets = paddle.concat( (-teacher_deltas[b_mask][..., :2], teacher_deltas[b_mask][..., 2:]), -1) iou_loss = GIoULoss(reduction='mean') loss_iou = iou_loss(inputs, targets) # distill_loss_dfl loss_dfl = F.cross_entropy( student_dfl[b_mask].reshape([-1, reg_ch]), teacher_dfl[b_mask].reshape([-1, reg_ch]), soft_label=True, reduction='mean') return { "distill_loss_cls": loss_cls, "distill_loss_iou": loss_iou, "distill_loss_dfl": loss_dfl, "distill_loss_contrast": loss_contrast, "fg_sum": fg_num, }