更换文档检测模型
This commit is contained in:
236
paddle_detection/ppdet/modeling/ssod/losses.py
Normal file
236
paddle_detection/ppdet/modeling/ssod/losses.py
Normal file
@@ -0,0 +1,236 @@
|
||||
# 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,
|
||||
}
|
||||
Reference in New Issue
Block a user