Browse Source

Fix bug #541 #542 (#545)

* fix #541 #542

* Update train.py

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
5.0
AlexWang1900 GitHub 4 years ago
parent
commit
a209a32019
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 3 deletions
  1. +1
    -1
      train.py
  2. +2
    -2
      utils/utils.py

+ 1
- 1
train.py View File



# Save last, best and delete # Save last, best and delete
torch.save(ckpt, last) torch.save(ckpt, last)
if (best_fitness == fi) and not final_epoch:
if best_fitness == fi:
torch.save(ckpt, best) torch.save(ckpt, best)
del ckpt del ckpt
# end epoch ---------------------------------------------------------------------------------------------------- # end epoch ----------------------------------------------------------------------------------------------------

+ 2
- 2
utils/utils.py View File

elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 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 / (1 - iou + v)
return iou - (rho2 / c2 + v * alpha) # CIoU
alpha = v / (1 - iou + v + 1e-16)
return iou - (rho2 / c2 + v * alpha ) # CIoU


return iou return iou



Loading…
Cancel
Save