Refactor Detect() anchors for ONNX <> OpenCV DNN compatibility (#4833)

* refactor anchors and anchor_grid in Detect Layer

* fix CI failures by adding compatibility

* fix tf failure

* fix different devices errors

* Cleanup

* fix anchors overwriting issue

* better refactoring

* Remove self.anchor_grid shape check (redundant with self.grid check)

Also PEP8 / 120 line width

* Convert _make_grid() from static to dynamic method

* Remove anchor_grid.to(device)

clone() should already clone to same device as self.anchors

* fix different devices error

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Jebastin NadarandGlenn Jocher authored and GitHub committed 2021-10-11 09:58:42 -07:00
1 parent 153873e9e4
commit 9d75e42f98
5 files changed
+24 -16

No files matched your search

+4 -6
View File
@@ -15,13 +15,12 @@ from utils.general import colorstr
def check_anchor_order(m):
# Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary
a = m.anchor_grid.prod(-1).view(-1) # anchor area
a = m.anchors.prod(-1).view(-1) # anchor area
da = a[-1] - a[0] # delta a
ds = m.stride[-1] - m.stride[0] # delta s
if da.sign() != ds.sign(): # same order
print('Reversing anchor order')
m.anchors[:] = m.anchors.flip(0)
m.anchor_grid[:] = m.anchor_grid.flip(0)
def check_anchors(dataset, model, thr=4.0, imgsz=640):
@@ -41,12 +40,12 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
bpr = (best > 1. / thr).float().mean() # best possible recall
return bpr, aat
anchors = m.anchor_grid.clone().cpu().view(-1, 2) # current anchors
bpr, aat = metric(anchors)
anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1) # current anchors
bpr, aat = metric(anchors.cpu().view(-1, 2))
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
na = m.anchors.numel() // 2 # number of anchors
try:
anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
except Exception as e:
@@ -54,7 +53,6 @@ def check_anchors(dataset, model, thr=4.0, imgsz=640):
new_bpr = metric(anchors)[0]
if new_bpr > bpr: # replace anchors
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.')