提交 a411c825 编写于 作者: L LDOUBLEV

add sum loss

上级 80561b15
......@@ -18,19 +18,60 @@ import paddle.nn as nn
from .rec_ctc_loss import CTCLoss
from .basic_loss import DMLLoss
from .basic_loss import DistanceLoss
from .det_db_loss import DBLoss
from .det_basic_loss import BalanceLoss, MaskL1Loss, DiceLoss
class DistillationDMLLoss(DMLLoss):
def _sum_loss(loss_dict):
if "loss" in loss_dict.keys():
return loss_dict
else:
loss_dict["loss"] = 0.
for k, value in loss_dict.items():
if k == "loss":
continue
else:
loss_dict["loss"] += value
return loss_dict
# class DistillationDMLLoss(DMLLoss):
"""
"""
def __init__(self, model_name_pairs=[], act=None, key=None,
def __init__(self,
model_name_pairs=[],
act=None,
key=None,
maps_name=None,
name="loss_dml"):
super().__init__(act=act)
assert isinstance(model_name_pairs, list)
self.key = key
self.model_name_pairs = model_name_pairs
self.name = name
self.maps_name = self.maps_name
def _check_maps_name(self, maps_name):
if maps_name is None:
return None
elif type(maps_name) == str:
return [maps_name]
elif type(maps_name) == list:
return [maps_name]
else:
return None
def _slice_out(self, outs):
new_outs = {}
for k in self.maps_name:
if k == "thrink_maps":
new_outs[k] = paddle.slice(outs, axes=1, starts=0, ends=1)
elif k == "threshold_maps":
new_outs[k] = paddle.slice(outs, axes=1, starts=1, ends=2)
elif k == "binary_maps":
new_outs[k] = paddle.slice(outs, axes=1, starts=2, ends=3)
else:
continue
def forward(self, predicts, batch):
loss_dict = dict()
......@@ -40,6 +81,8 @@ class DistillationDMLLoss(DMLLoss):
if self.key is not None:
out1 = out1[self.key]
out2 = out2[self.key]
if self.maps_name is None:
loss = super().forward(out1, out2)
if isinstance(loss, dict):
for key in loss:
......@@ -47,6 +90,21 @@ class DistillationDMLLoss(DMLLoss):
idx)] = loss[key]
else:
loss_dict["{}_{}".format(self.name, idx)] = loss
else:
outs1 = self._slice_out(out1)
outs2 = self._slice_out(out2)
for k in outs1.keys():
loss = super().forward(outs1[k], outs2[k])
if isinstance(loss, dict):
for key in loss:
loss_dict["{}_{}_{}_{}_{}".format(key, pair[
0], pair[1], map_name, idx)] = loss[key]
else:
loss_dict["{}_{}_{}".format(self.name, map_name,
idx)] = loss
loss_dict = _sum_loss(loss_dict)
return loss_dict
......@@ -73,6 +131,93 @@ class DistillationCTCLoss(CTCLoss):
return loss_dict
"""
class DistillationDBLoss(DBLoss):
def __init__(self,
model_name_list=[],
balance_loss=True,
main_loss_type='DiceLoss',
alpha=5,
beta=10,
ohem_ratio=3,
eps=1e-6,
name="db_loss",
**kwargs):
super().__init__()
self.model_name_list = model_name_list
self.name = name
def forward(self, predicts, batch):
loss_dict = dict()
for idx, model_name in enumerate(self.model_name_list):
out = predicts[model_name]
if self.key is not None:
out = out[self.key]
loss = super().forward(out, batch)
if isinstance(loss, dict):
for key in loss.keys():
if key == "loss":
continue
loss_dict[f"{self.name}_{model_name}_{key}"] = loss[key]
else:
loss_dict[f"{self.name}_{model_name}"] = loss
loss_dict = _sum_loss(loss_dict)
return loss_dict
class DistillationDilaDBLoss(DBLoss):
def __init__(self, model_name_pairs=[],
balance_loss=True,
main_loss_type='DiceLoss',
alpha=5,
beta=10,
ohem_ratio=3,
eps=1e-6,
name="dila_dbloss"):
super().__init__()
self.model_name_pairs = model_name_pairs
self.name = name
def forward(self, predicts, batch):
loss_dict = dict()
for idx, pair in enumerate(self.model_name_pairs):
stu_outs = predicts[pair[0]]
tch_outs = predicts[pair[1]]
if self.key is not None:
stu_preds = stu_outs[self.key]
tch_preds = tch_outs[self.key]
stu_shrink_maps = stu_preds[:, 0, :, :]
stu_binary_maps = stu_preds[:, 2, :, :]
# dilation to teacher prediction
dilation_w = np.array([[1,1], [1,1]])
th_shrink_maps = tch_preds[:, 0, :, :]
th_shrink_maps = th_shrink_maps.numpy() > 0.3 # thresh = 0.3
dilate_maps = np.zeros_like(th_shrink_maps).astype(np.float32)
for i in range(th_shrink_maps.shape[0]):
dilate_maps[i] = cv2.dilate(th_shrink_maps[i, :, :].astype(np.uint8), dilation_w)
th_shrink_maps = paddle.to_tensor(dilate_maps)
label_threshold_map, label_threshold_mask, label_shrink_map, label_shrink_mask = batch[1:]
# calculate the shrink map loss
bce_loss = self.alpha * self.bce_loss(stu_shrink_maps, th_shrink_maps,
label_shrink_mask)
loss_binary_maps = self.dice_loss(stu_binary_maps, th_shrink_maps,
label_shrink_mask)
k = f"{self.name}_{pair[0]}_{pair[1]}"
loss_dict[k] = bce_loss + loss_binary_maps
loss_dict = _sum_loss(loss_dict)
return loss
"""
class DistillationDistanceLoss(DistanceLoss):
"""
"""
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册