* fix #541 #542 * Update train.py Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>5.0
# 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 ---------------------------------------------------------------------------------------------------- |
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 | ||||