apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update loss.py (#1959)

Browse Source 
* Update loss.py

* Update metrics.py

* Update loss.py
This commit is contained in:
Glenn Jocher 2022-06-29 18:07:06 +02:00 committed by GitHub
parent 0aa65efcc0
commit 7ec9614961
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 70 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

88
utils/loss.py
View File

@ -7,7 +7,7 @@ import torch
import torch.nn as nn import torch.nn as nn
from utils.metrics import bbox_iou from utils.metrics import bbox_iou
from utils.torch_utils import is_parallel from utils.torch_utils import de_parallel
def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441 def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
@ -89,9 +89,10 @@ class QFocalLoss(nn.Module):
class ComputeLoss: class ComputeLoss:
sort_obj_iou = False
# Compute losses # Compute losses
def __init__(self, model, autobalance=False): def __init__(self, model, autobalance=False):
self.sort_obj_iou = False
device = next(model.parameters()).device # get model device device = next(model.parameters()).device # get model device
h = model.hyp # hyperparameters h = model.hyp # hyperparameters
@ -107,46 +108,53 @@ class ComputeLoss:
if g > 0: if g > 0:
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module m = de_parallel(model).model[-1] # Detect() module
self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7 self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7
self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index self.ssi = list(m.stride).index(16) if autobalance else 0 # stride 16 index
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
for k in 'na', 'nc', 'nl', 'anchors': self.na = m.na # number of anchors
setattr(self, k, getattr(det, k)) self.nc = m.nc # number of classes
self.nl = m.nl # number of layers
self.anchors = m.anchors
self.device = device
def __call__(self, p, targets): # predictions, targets, model def __call__(self, p, targets): # predictions, targets
device = targets.device lcls = torch.zeros(1, device=self.device) # class loss
lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device) lbox = torch.zeros(1, device=self.device) # box loss
lobj = torch.zeros(1, device=self.device) # object loss
tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets
# Losses # Losses
for i, pi in enumerate(p): # layer index, layer predictions for i, pi in enumerate(p): # layer index, layer predictions
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
tobj = torch.zeros_like(pi[..., 0], device=device) # target obj tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device) # target obj
n = b.shape[0] # number of targets n = b.shape[0] # number of targets
if n: if n:
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets # pxy, pwh, _, pcls = pi[b, a, gj, gi].tensor_split((2, 4, 5), dim=1) # faster, requires torch 1.8.0
pxy, pwh, _, pcls = pi[b, a, gj, gi].split((2, 2, 1, self.nc), 1) # target-subset of predictions
# Regression # Regression
pxy = ps[:, :2].sigmoid() * 2 - 0.5 pxy = pxy.sigmoid() * 2 - 0.5
pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i] pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
pbox = torch.cat((pxy, pwh), 1) # predicted box pbox = torch.cat((pxy, pwh), 1) # predicted box
iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target) iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze() # iou(prediction, target)
lbox += (1.0 - iou).mean() # iou loss lbox += (1.0 - iou).mean() # iou loss
# Objectness # Objectness
score_iou = iou.detach().clamp(0).type(tobj.dtype) iou = iou.detach().clamp(0).type(tobj.dtype)
if self.sort_obj_iou: if self.sort_obj_iou:
sort_id = torch.argsort(score_iou) j = iou.argsort()
b, a, gj, gi, score_iou = b[sort_id], a[sort_id], gj[sort_id], gi[sort_id], score_iou[sort_id] b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou # iou ratio if self.gr < 1:
iou = (1.0 - self.gr) + self.gr * iou
tobj[b, a, gj, gi] = iou # iou ratio
# Classification # Classification
if self.nc > 1: # cls loss (only if multiple classes) if self.nc > 1: # cls loss (only if multiple classes)
t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets t = torch.full_like(pcls, self.cn, device=self.device) # targets
t[range(n), tcls[i]] = self.cp t[range(n), tcls[i]] = self.cp
lcls += self.BCEcls(ps[:, 5:], t) # BCE lcls += self.BCEcls(pcls, t) # BCE
# Append targets to text file # Append targets to text file
# with open('targets.txt', 'a') as file: # with open('targets.txt', 'a') as file:
@ -170,25 +178,31 @@ class ComputeLoss:
# Build targets for compute_loss(), input targets(image,class,x,y,w,h) # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
na, nt = self.na, targets.shape[0] # number of anchors, targets na, nt = self.na, targets.shape[0] # number of anchors, targets
tcls, tbox, indices, anch = [], [], [], [] tcls, tbox, indices, anch = [], [], [], []
gain = torch.ones(7, device=targets.device) # normalized to gridspace gain gain = torch.ones(7, device=self.device) # normalized to gridspace gain
ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt) ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)
targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2) # append anchor indices
g = 0.5 # bias g = 0.5 # bias
off = torch.tensor([[0, 0], off = torch.tensor(
[1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m [
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm [0, 0],
], device=targets.device).float() * g # offsets [1, 0],
[0, 1],
[-1, 0],
[0, -1], # j,k,l,m
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
],
device=self.device).float() * g # offsets
for i in range(self.nl): for i in range(self.nl):
anchors = self.anchors[i] anchors, shape = self.anchors[i], p[i].shape
gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]] # xyxy gain
# Match targets to anchors # Match targets to anchors
t = targets * gain t = targets * gain # shape(3,n,7)
if nt: if nt:
# Matches # Matches
r = t[:, :, 4:6] / anchors[:, None] # wh ratio r = t[..., 4:6] / anchors[:, None] # wh ratio
j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare
# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2)) # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
t = t[j] # filter t = t[j] # filter
@ -206,15 +220,13 @@ class ComputeLoss:
offsets = 0 offsets = 0
# Define # Define
b, c = t[:, :2].long().T # image, class bc, gxy, gwh, a = t.chunk(4, 1) # (image, class), grid xy, grid wh, anchors
gxy = t[:, 2:4] # grid xy a, (b, c) = a.long().view(-1), bc.long().T # anchors, image, class
gwh = t[:, 4:6] # grid wh
gij = (gxy - offsets).long() gij = (gxy - offsets).long()
gi, gj = gij.T # grid xy indices gi, gj = gij.T # grid indices
# Append # Append
a = t[:, 6].long() # anchor indices indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1))) # image, anchor, grid
indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
anch.append(anchors[a]) # anchors anch.append(anchors[a]) # anchors
tcls.append(c) # class tcls.append(c) # class

44
utils/metrics.py
View File

@ -189,49 +189,45 @@ class ConfusionMatrix:
print(' '.join(map(str, self.matrix[i]))) print(' '.join(map(str, self.matrix[i])))
def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7): def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
# Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 # Returns Intersection over Union (IoU) of box1(1,4) to box2(n,4)
box2 = box2.T
# Get the coordinates of bounding boxes # Get the coordinates of bounding boxes
if x1y1x2y2: # x1, y1, x2, y2 = box1 if xywh: # transform from xywh to xyxy
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3] (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, 1), box2.chunk(4, 1)
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
else: # transform from xywh to xyxy b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2 b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2 else: # x1, y1, x2, y2 = box1
b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2 b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, 1)
b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2 b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, 1)
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
# Intersection area # Intersection area
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
# Union Area # Union Area
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
union = w1 * h1 + w2 * h2 - inter + eps union = w1 * h1 + w2 * h2 - inter + eps
# IoU
iou = inter / union iou = inter / union
if GIoU or DIoU or CIoU: if CIoU or DIoU or GIoU:
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center dist ** 2
(b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared if CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
if DIoU:
return iou - rho2 / c2 # DIoU
elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
with torch.no_grad(): with torch.no_grad():
alpha = v / (v - iou + (1 + eps)) alpha = v / (v - iou + (1 + eps))
return iou - (rho2 / c2 + v * alpha) # CIoU return iou - (rho2 / c2 + v * alpha) # CIoU
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf return iou - rho2 / c2 # DIoU
c_area = cw * ch + eps # convex area c_area = cw * ch + eps # convex area
return iou - (c_area - union) / c_area # GIoU return iou - (c_area - union) / c_area # GIoU https://arxiv.org/pdf/1902.09630.pdf
else: return iou # IoU
return iou # IoU
def box_iou(box1, box2): def box_iou(box1, box2):