This commit is contained in:
Glenn Jocher 2019-08-18 14:20:46 +02:00
parent 7ee28a7bb6
commit 4050650669
2 changed files with 34 additions and 22 deletions

View File

@ -155,14 +155,17 @@ class YOLOLayer(nn.Module):
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method # io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
io[..., :4] *= self.stride io[..., :4] *= self.stride
arc = 'normal' # (normal, uCE, uBCE) architecture types arc = 'normal' # (normal, uCE, uBCE, uBCEs) detection architectures
if arc == 'normal': if arc == 'normal':
io[..., 4:] = torch.sigmoid(io[..., 4:]) torch.sigmoid_(io[..., 4:])
elif arc == 'uCE': elif arc == 'uCE': # unified CE (1 background + 80 classes)
io[..., 4:] = F.softmax(io[..., 4:], dim=4) # unified detection CE io[..., 4:] = F.softmax(io[..., 4:], dim=4)
io[..., 4] = 1 io[..., 4] = 1
elif arc == 'uBCE': elif arc == 'uBCE': # unified BCE (1 background + 80 classes)
io[..., 4] = io[..., 5:].max(4)[0] # unified detection BCE torch.sigmoid_(io[..., 4:])
io[..., 4] = 1 - io[..., 4]
elif arc == 'uBCEs': # unified BCE simplified (80 classes)
torch.sigmoid_(io[..., 4:])
if self.nc == 1: if self.nc == 1:
io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235 io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235

View File

@ -321,12 +321,12 @@ def compute_loss(p, targets, model): # predictions, targets, model
# Define criteria # Define criteria
BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']])) BCEcls = nn.BCEWithLogitsLoss(pos_weight=ft([h['cls_pw']]))
BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']])) BCEobj = nn.BCEWithLogitsLoss(pos_weight=ft([h['obj_pw']]))
CE = nn.CrossEntropyLoss(weight=model.class_weights) # CE = nn.CrossEntropyLoss(weight=model.class_weights)
# Compute losses # Compute losses
bs = p[0].shape[0] # batch size bs = p[0].shape[0] # batch size
k = bs / 64 # loss gain k = bs / 64 # loss gain
arc = 'normal' # (normal, uCE, uBCE) architecture types arc = 'normal' # (normal, uCE, uBCE, uBCEs) detection architectures
for i, pi0 in enumerate(p): # layer i predictions, i for i, pi0 in enumerate(p): # layer i predictions, i
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
tobj = torch.zeros_like(pi0[..., 0]) # target obj tobj = torch.zeros_like(pi0[..., 0]) # target obj
@ -342,33 +342,42 @@ def compute_loss(p, targets, model): # predictions, targets, model
pxy = torch.sigmoid(pi[..., 0:2]) # pxy = pxy * s - (s - 1) / 2, s = 1.5 (scale_xy) pxy = torch.sigmoid(pi[..., 0:2]) # pxy = pxy * s - (s - 1) / 2, s = 1.5 (scale_xy)
pbox = torch.cat((pxy, torch.exp(pi[..., 2:4]) * anchor_vec[i]), 1) # predicted pbox = torch.cat((pxy, torch.exp(pi[..., 2:4]) * anchor_vec[i]), 1) # predicted
giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True) # giou computation
lbox += (k * h['giou']) * (1.0 - giou).mean() # giou loss lbox += (1.0 - giou).mean() # giou loss
if arc == 'normal' and model.nc > 1: # cls loss (only if multiple classes) if arc == 'normal' and model.nc > 1: # cls loss (only if multiple classes)
tclsm = torch.zeros_like(pi[..., 5:]) t = torch.zeros_like(pi[..., 5:]) # targets
tclsm[range(nb), tcls[i]] = 1.0 t[range(nb), tcls[i]] = 1.0
lcls += (k * h['cls']) * BCEcls(pi[..., 5:], tclsm) # BCE lcls += BCEcls(pi[..., 5:], t) # BCE
# lcls += (k * h['cls']) * CE(pi[..., 5:], tcls[i]) # CE # lcls += CE(pi[..., 5:], tcls[i]) # CE
# Append targets to text file # Append targets to text file
# with open('targets.txt', 'a') as file: # with open('targets.txt', 'a') as file:
# [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)] # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
if arc == 'normal': if arc == 'normal':
lobj += (k * h['obj']) * BCEobj(pi0[..., 4], tobj) # obj loss lobj += BCEobj(pi0[..., 4], tobj) # obj loss
elif arc == 'uCE': # suggest h['cls']=5. elif arc == 'uCE': # unified CE (1 background + 80 classes), hyps 20
udm_ce = torch.zeros_like(pi0[..., 0]).long() # unified detection matrix for CE t = torch.zeros_like(pi0[..., 0], dtype=torch.long) # targets
if nb: if nb:
udm_ce[b, a, gj, gi] = tcls[i] + 1 t[b, a, gj, gi] = tcls[i] + 1
lcls += (k * h['cls']) * CE(pi0[..., 4:].view(-1, model.nc + 1), udm_ce.view(-1)) # unified CE lcls += CE(pi0[..., 4:].view(-1, model.nc + 1), t.view(-1))
elif arc == 'uBCE': elif arc == 'uBCE': # unified BCE (1 background + 80 classes), hyps 200-30
udm = torch.zeros_like(pi0[..., 5:]) # unified detection matrix for BCE t = torch.zeros_like(pi0[..., 5:]) # targets
if nb: if nb:
udm[b, a, gj, gi, tcls[i]] = 1.0 t[b, a, gj, gi, tcls[i]] = 1.0
lcls += (k * h['cls']) * BCEcls(pi0[..., 5:], udm) # unified BCE (hyps 200-30) lcls += BCEcls(pi0[..., 5:], t)
elif arc == 'uBCEs': # unified BCE simplified (80 classes)
t = torch.zeros_like(pi0[..., 5:]) # targets
if nb:
t[b, a, gj, gi, tcls[i]] = 1.0
lcls += BCEcls(pi0[..., 5:], t)
lbox *= k * h['giou']
lobj *= k * h['obj']
lcls *= k * h['cls']
loss = lbox + lobj + lcls loss = lbox + lobj + lcls
return loss, torch.cat((lbox, ft([0]), lobj, lcls, loss)).detach() return loss, torch.cat((lbox, ft([0]), lobj, lcls, loss)).detach()