Scipy kmeans-robust autoanchor update (#2470)
Fix for https://github.com/ultralytics/yolov5/issues/2394
This commit is contained in:
1 parent
38ff499b26
commit
2d41e70e82
1 file changed
+11
-6
+11
-6
@@ -37,17 +37,21 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
|
||||
bpr = (best > 1. / thr).float().mean() # best possible recall
|
||||
return bpr, aat
|
||||
|
||||
bpr, aat = metric(m.anchor_grid.clone().cpu().view(-1, 2))
|
||||
anchors = m.anchor_grid.clone().cpu().view(-1, 2) # current anchors
|
||||
bpr, aat = metric(anchors)
|
||||
print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='')
|
||||
if bpr < 0.98: # threshold to recompute
|
||||
print('. Attempting to improve anchors, please wait...')
|
||||
na = m.anchor_grid.numel() // 2 # number of anchors
|
||||
new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
|
||||
new_bpr = metric(new_anchors.reshape(-1, 2))[0]
|
||||
try:
|
||||
anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
|
||||
except Exception as e:
|
||||
print(f'{prefix}ERROR: {e}')
|
||||
new_bpr = metric(anchors)[0]
|
||||
if new_bpr > bpr: # replace anchors
|
||||
new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors)
|
||||
m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid) # for inference
|
||||
m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
|
||||
anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
|
||||
m.anchor_grid[:] = anchors.clone().view_as(m.anchor_grid) # for inference
|
||||
m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
|
||||
check_anchor_order(m)
|
||||
print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
|
||||
else:
|
||||
@@ -119,6 +123,7 @@ def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=10
|
||||
print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...')
|
||||
s = wh.std(0) # sigmas for whitening
|
||||
k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
|
||||
assert len(k) == n, print(f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}')
|
||||
k *= s
|
||||
wh = torch.tensor(wh, dtype=torch.float32) # filtered
|
||||
wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
|
||||
|
||||
Reference in new issue
Block a user