|
| 1 | +# Auto-anchor utils |
| 2 | + |
| 3 | +import numpy as np |
| 4 | +import torch |
| 5 | +import yaml |
| 6 | +from scipy.cluster.vq import kmeans |
| 7 | +from tqdm import tqdm |
| 8 | + |
| 9 | +from utils.general import colorstr |
| 10 | + |
| 11 | + |
| 12 | +def check_anchor_order(m): |
| 13 | + # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary |
| 14 | + a = m.anchor_grid.prod(-1).view(-1) # anchor area |
| 15 | + da = a[-1] - a[0] # delta a |
| 16 | + ds = m.stride[-1] - m.stride[0] # delta s |
| 17 | + if da.sign() != ds.sign(): # same order |
| 18 | + print('Reversing anchor order') |
| 19 | + m.anchors[:] = m.anchors.flip(0) |
| 20 | + m.anchor_grid[:] = m.anchor_grid.flip(0) |
| 21 | + |
| 22 | + |
| 23 | +def check_anchors(dataset, model, thr=4.0, imgsz=640): |
| 24 | + # Check anchor fit to data, recompute if necessary |
| 25 | + prefix = colorstr('autoanchor: ') |
| 26 | + print(f'\n{prefix}Analyzing anchors... ', end='') |
| 27 | + m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() |
| 28 | + shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) |
| 29 | + scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale |
| 30 | + wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh |
| 31 | + |
| 32 | + def metric(k): # compute metric |
| 33 | + r = wh[:, None] / k[None] |
| 34 | + x = torch.min(r, 1. / r).min(2)[0] # ratio metric |
| 35 | + best = x.max(1)[0] # best_x |
| 36 | + aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold |
| 37 | + bpr = (best > 1. / thr).float().mean() # best possible recall |
| 38 | + return bpr, aat |
| 39 | + |
| 40 | + bpr, aat = metric(m.anchor_grid.clone().cpu().view(-1, 2)) |
| 41 | + print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='') |
| 42 | + if bpr < 0.98: # threshold to recompute |
| 43 | + print('. Attempting to improve anchors, please wait...') |
| 44 | + na = m.anchor_grid.numel() // 2 # number of anchors |
| 45 | + new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) |
| 46 | + new_bpr = metric(new_anchors.reshape(-1, 2))[0] |
| 47 | + if new_bpr > bpr: # replace anchors |
| 48 | + new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors) |
| 49 | + m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid) # for inference |
| 50 | + m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss |
| 51 | + check_anchor_order(m) |
| 52 | + print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.') |
| 53 | + else: |
| 54 | + print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.') |
| 55 | + print('') # newline |
| 56 | + |
| 57 | + |
| 58 | +def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): |
| 59 | + """ Creates kmeans-evolved anchors from training dataset |
| 60 | +
|
| 61 | + Arguments: |
| 62 | + path: path to dataset *.yaml, or a loaded dataset |
| 63 | + n: number of anchors |
| 64 | + img_size: image size used for training |
| 65 | + thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 |
| 66 | + gen: generations to evolve anchors using genetic algorithm |
| 67 | + verbose: print all results |
| 68 | +
|
| 69 | + Return: |
| 70 | + k: kmeans evolved anchors |
| 71 | +
|
| 72 | + Usage: |
| 73 | + from utils.autoanchor import *; _ = kmean_anchors() |
| 74 | + """ |
| 75 | + thr = 1. / thr |
| 76 | + prefix = colorstr('autoanchor: ') |
| 77 | + |
| 78 | + def metric(k, wh): # compute metrics |
| 79 | + r = wh[:, None] / k[None] |
| 80 | + x = torch.min(r, 1. / r).min(2)[0] # ratio metric |
| 81 | + # x = wh_iou(wh, torch.tensor(k)) # iou metric |
| 82 | + return x, x.max(1)[0] # x, best_x |
| 83 | + |
| 84 | + def anchor_fitness(k): # mutation fitness |
| 85 | + _, best = metric(torch.tensor(k, dtype=torch.float32), wh) |
| 86 | + return (best * (best > thr).float()).mean() # fitness |
| 87 | + |
| 88 | + def print_results(k): |
| 89 | + k = k[np.argsort(k.prod(1))] # sort small to large |
| 90 | + x, best = metric(k, wh0) |
| 91 | + bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr |
| 92 | + print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr') |
| 93 | + print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' |
| 94 | + f'past_thr={x[x > thr].mean():.3f}-mean: ', end='') |
| 95 | + for i, x in enumerate(k): |
| 96 | + print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg |
| 97 | + return k |
| 98 | + |
| 99 | + if isinstance(path, str): # *.yaml file |
| 100 | + with open(path) as f: |
| 101 | + data_dict = yaml.load(f, Loader=yaml.SafeLoader) # model dict |
| 102 | + from utils.datasets import LoadImagesAndLabels |
| 103 | + dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) |
| 104 | + else: |
| 105 | + dataset = path # dataset |
| 106 | + |
| 107 | + # Get label wh |
| 108 | + shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) |
| 109 | + wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh |
| 110 | + |
| 111 | + # Filter |
| 112 | + i = (wh0 < 3.0).any(1).sum() |
| 113 | + if i: |
| 114 | + print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.') |
| 115 | + wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels |
| 116 | + # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1 |
| 117 | + |
| 118 | + # Kmeans calculation |
| 119 | + print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...') |
| 120 | + s = wh.std(0) # sigmas for whitening |
| 121 | + k, dist = kmeans(wh / s, n, iter=30) # points, mean distance |
| 122 | + k *= s |
| 123 | + wh = torch.tensor(wh, dtype=torch.float32) # filtered |
| 124 | + wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered |
| 125 | + k = print_results(k) |
| 126 | + |
| 127 | + # Plot |
| 128 | + # k, d = [None] * 20, [None] * 20 |
| 129 | + # for i in tqdm(range(1, 21)): |
| 130 | + # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance |
| 131 | + # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True) |
| 132 | + # ax = ax.ravel() |
| 133 | + # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') |
| 134 | + # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh |
| 135 | + # ax[0].hist(wh[wh[:, 0]<100, 0],400) |
| 136 | + # ax[1].hist(wh[wh[:, 1]<100, 1],400) |
| 137 | + # fig.savefig('wh.png', dpi=200) |
| 138 | + |
| 139 | + # Evolve |
| 140 | + npr = np.random |
| 141 | + f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma |
| 142 | + pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:') # progress bar |
| 143 | + for _ in pbar: |
| 144 | + v = np.ones(sh) |
| 145 | + while (v == 1).all(): # mutate until a change occurs (prevent duplicates) |
| 146 | + v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) |
| 147 | + kg = (k.copy() * v).clip(min=2.0) |
| 148 | + fg = anchor_fitness(kg) |
| 149 | + if fg > f: |
| 150 | + f, k = fg, kg.copy() |
| 151 | + pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}' |
| 152 | + if verbose: |
| 153 | + print_results(k) |
| 154 | + |
| 155 | + return print_results(k) |
0 commit comments