123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185 |
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- from ..general import xywh2xyxy
- from ..loss import FocalLoss, smooth_BCE
- from ..metrics import bbox_iou
- from ..torch_utils import de_parallel
- from .general import crop_mask
- class ComputeLoss:
-
- def __init__(self, model, autobalance=False, overlap=False):
- self.sort_obj_iou = False
- self.overlap = overlap
- device = next(model.parameters()).device
- h = model.hyp
-
- BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
- BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
-
- self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0))
-
- g = h['fl_gamma']
- if g > 0:
- BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
- m = de_parallel(model).model[-1]
- self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02])
- self.ssi = list(m.stride).index(16) if autobalance else 0
- self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
- self.na = m.na
- self.nc = m.nc
- self.nl = m.nl
- self.nm = m.nm
- self.anchors = m.anchors
- self.device = device
- def __call__(self, preds, targets, masks):
- p, proto = preds
- bs, nm, mask_h, mask_w = proto.shape
- lcls = torch.zeros(1, device=self.device)
- lbox = torch.zeros(1, device=self.device)
- lobj = torch.zeros(1, device=self.device)
- lseg = torch.zeros(1, device=self.device)
- tcls, tbox, indices, anchors, tidxs, xywhn = self.build_targets(p, targets)
-
- for i, pi in enumerate(p):
- b, a, gj, gi = indices[i]
- tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device)
- n = b.shape[0]
- if n:
- pxy, pwh, _, pcls, pmask = pi[b, a, gj, gi].split((2, 2, 1, self.nc, nm), 1)
-
- pxy = pxy.sigmoid() * 2 - 0.5
- pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
- pbox = torch.cat((pxy, pwh), 1)
- iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze()
- lbox += (1.0 - iou).mean()
-
- iou = iou.detach().clamp(0).type(tobj.dtype)
- if self.sort_obj_iou:
- j = iou.argsort()
- b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
- if self.gr < 1:
- iou = (1.0 - self.gr) + self.gr * iou
- tobj[b, a, gj, gi] = iou
-
- if self.nc > 1:
- t = torch.full_like(pcls, self.cn, device=self.device)
- t[range(n), tcls[i]] = self.cp
- lcls += self.BCEcls(pcls, t)
-
- if tuple(masks.shape[-2:]) != (mask_h, mask_w):
- masks = F.interpolate(masks[None], (mask_h, mask_w), mode='nearest')[0]
- marea = xywhn[i][:, 2:].prod(1)
- mxyxy = xywh2xyxy(xywhn[i] * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device))
- for bi in b.unique():
- j = b == bi
- if self.overlap:
- mask_gti = torch.where(masks[bi][None] == tidxs[i][j].view(-1, 1, 1), 1.0, 0.0)
- else:
- mask_gti = masks[tidxs[i]][j]
- lseg += self.single_mask_loss(mask_gti, pmask[j], proto[bi], mxyxy[j], marea[j])
- obji = self.BCEobj(pi[..., 4], tobj)
- lobj += obji * self.balance[i]
- if self.autobalance:
- self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
- if self.autobalance:
- self.balance = [x / self.balance[self.ssi] for x in self.balance]
- lbox *= self.hyp['box']
- lobj *= self.hyp['obj']
- lcls *= self.hyp['cls']
- lseg *= self.hyp['box'] / bs
- loss = lbox + lobj + lcls + lseg
- return loss * bs, torch.cat((lbox, lseg, lobj, lcls)).detach()
- def single_mask_loss(self, gt_mask, pred, proto, xyxy, area):
-
- pred_mask = (pred @ proto.view(self.nm, -1)).view(-1, *proto.shape[1:])
- loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction='none')
- return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).mean()
- def build_targets(self, p, targets):
-
- na, nt = self.na, targets.shape[0]
- tcls, tbox, indices, anch, tidxs, xywhn = [], [], [], [], [], []
- gain = torch.ones(8, device=self.device)
- ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt)
- if self.overlap:
- batch = p[0].shape[0]
- ti = []
- for i in range(batch):
- num = (targets[:, 0] == i).sum()
- ti.append(torch.arange(num, device=self.device).float().view(1, num).repeat(na, 1) + 1)
- ti = torch.cat(ti, 1)
- else:
- ti = torch.arange(nt, device=self.device).float().view(1, nt).repeat(na, 1)
- targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None], ti[..., None]), 2)
- g = 0.5
- off = torch.tensor(
- [
- [0, 0],
- [1, 0],
- [0, 1],
- [-1, 0],
- [0, -1],
-
- ],
- device=self.device).float() * g
- for i in range(self.nl):
- anchors, shape = self.anchors[i], p[i].shape
- gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]
-
- t = targets * gain
- if nt:
-
- r = t[..., 4:6] / anchors[:, None]
- j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']
-
- t = t[j]
-
- gxy = t[:, 2:4]
- gxi = gain[[2, 3]] - gxy
- j, k = ((gxy % 1 < g) & (gxy > 1)).T
- l, m = ((gxi % 1 < g) & (gxi > 1)).T
- j = torch.stack((torch.ones_like(j), j, k, l, m))
- t = t.repeat((5, 1, 1))[j]
- offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
- else:
- t = targets[0]
- offsets = 0
-
- bc, gxy, gwh, at = t.chunk(4, 1)
- (a, tidx), (b, c) = at.long().T, bc.long().T
- gij = (gxy - offsets).long()
- gi, gj = gij.T
-
- indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))
- tbox.append(torch.cat((gxy - gij, gwh), 1))
- anch.append(anchors[a])
- tcls.append(c)
- tidxs.append(tidx)
- xywhn.append(torch.cat((gxy, gwh), 1) / gain[2:6])
- return tcls, tbox, indices, anch, tidxs, xywhn
|