You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

169 lines
7.2KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. AutoAnchor utils
  4. """
  5. import random
  6. import numpy as np
  7. import torch
  8. import yaml
  9. from tqdm import tqdm
  10. from utils.general import LOGGER, colorstr, emojis
  11. PREFIX = colorstr('AutoAnchor: ')
  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.anchors.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. LOGGER.info(f'{PREFIX}Reversing anchor order')
  19. m.anchors[:] = m.anchors.flip(0)
  20. def check_anchors(dataset, model, thr=4.0, imgsz=640):
  21. # Check anchor fit to data, recompute if necessary
  22. m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect()
  23. shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
  24. scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale
  25. wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh
  26. def metric(k): # compute metric
  27. r = wh[:, None] / k[None]
  28. x = torch.min(r, 1 / r).min(2)[0] # ratio metric
  29. best = x.max(1)[0] # best_x
  30. aat = (x > 1 / thr).float().sum(1).mean() # anchors above threshold
  31. bpr = (best > 1 / thr).float().mean() # best possible recall
  32. return bpr, aat
  33. anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1) # current anchors
  34. bpr, aat = metric(anchors.cpu().view(-1, 2))
  35. s = f'\n{PREFIX}{aat:.2f} anchors/target, {bpr:.3f} Best Possible Recall (BPR). '
  36. if bpr > 0.98: # threshold to recompute
  37. LOGGER.info(emojis(f'{s}Current anchors are a good fit to dataset ✅'))
  38. else:
  39. LOGGER.info(emojis(f'{s}Anchors are a poor fit to dataset ⚠️, attempting to improve...'))
  40. na = m.anchors.numel() // 2 # number of anchors
  41. try:
  42. anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
  43. except Exception as e:
  44. LOGGER.info(f'{PREFIX}ERROR: {e}')
  45. new_bpr = metric(anchors)[0]
  46. if new_bpr > bpr: # replace anchors
  47. anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
  48. m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
  49. check_anchor_order(m)
  50. s = f'{PREFIX}Done ✅ (optional: update model *.yaml to use these anchors in the future)'
  51. else:
  52. s = f'{PREFIX}Done ⚠️ (original anchors better than new anchors, proceeding with original anchors)'
  53. LOGGER.info(emojis(s))
  54. def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
  55. """ Creates kmeans-evolved anchors from training dataset
  56. Arguments:
  57. dataset: path to data.yaml, or a loaded dataset
  58. n: number of anchors
  59. img_size: image size used for training
  60. thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
  61. gen: generations to evolve anchors using genetic algorithm
  62. verbose: print all results
  63. Return:
  64. k: kmeans evolved anchors
  65. Usage:
  66. from utils.autoanchor import *; _ = kmean_anchors()
  67. """
  68. from scipy.cluster.vq import kmeans
  69. npr = np.random
  70. thr = 1 / thr
  71. def metric(k, wh): # compute metrics
  72. r = wh[:, None] / k[None]
  73. x = torch.min(r, 1 / r).min(2)[0] # ratio metric
  74. # x = wh_iou(wh, torch.tensor(k)) # iou metric
  75. return x, x.max(1)[0] # x, best_x
  76. def anchor_fitness(k): # mutation fitness
  77. _, best = metric(torch.tensor(k, dtype=torch.float32), wh)
  78. return (best * (best > thr).float()).mean() # fitness
  79. def print_results(k, verbose=True):
  80. k = k[np.argsort(k.prod(1))] # sort small to large
  81. x, best = metric(k, wh0)
  82. bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr
  83. s = f'{PREFIX}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr\n' \
  84. f'{PREFIX}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' \
  85. f'past_thr={x[x > thr].mean():.3f}-mean: '
  86. for i, x in enumerate(k):
  87. s += '%i,%i, ' % (round(x[0]), round(x[1]))
  88. if verbose:
  89. LOGGER.info(s[:-2])
  90. return k
  91. if isinstance(dataset, str): # *.yaml file
  92. with open(dataset, errors='ignore') as f:
  93. data_dict = yaml.safe_load(f) # model dict
  94. from utils.datasets import LoadImagesAndLabels
  95. dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
  96. # Get label wh
  97. shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
  98. wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh
  99. # Filter
  100. i = (wh0 < 3.0).any(1).sum()
  101. if i:
  102. LOGGER.info(f'{PREFIX}WARNING: Extremely small objects found: {i} of {len(wh0)} labels are < 3 pixels in size')
  103. wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
  104. # wh = wh * (npr.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
  105. # Kmeans init
  106. try:
  107. LOGGER.info(f'{PREFIX}Running kmeans for {n} anchors on {len(wh)} points...')
  108. assert n <= len(wh) # apply overdetermined constraint
  109. s = wh.std(0) # sigmas for whitening
  110. k = kmeans(wh / s, n, iter=30)[0] * s # points
  111. assert n == len(k) # kmeans may return fewer points than requested if wh is insufficient or too similar
  112. except Exception:
  113. LOGGER.warning(f'{PREFIX}WARNING: switching strategies from kmeans to random init')
  114. k = np.sort(npr.rand(n * 2)).reshape(n, 2) * img_size # random init
  115. wh, wh0 = (torch.tensor(x, dtype=torch.float32) for x in (wh, wh0))
  116. k = print_results(k, verbose=False)
  117. # Plot
  118. # k, d = [None] * 20, [None] * 20
  119. # for i in tqdm(range(1, 21)):
  120. # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance
  121. # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True)
  122. # ax = ax.ravel()
  123. # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
  124. # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh
  125. # ax[0].hist(wh[wh[:, 0]<100, 0],400)
  126. # ax[1].hist(wh[wh[:, 1]<100, 1],400)
  127. # fig.savefig('wh.png', dpi=200)
  128. # Evolve
  129. f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
  130. pbar = tqdm(range(gen), bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
  131. for _ in pbar:
  132. v = np.ones(sh)
  133. while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
  134. v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
  135. kg = (k.copy() * v).clip(min=2.0)
  136. fg = anchor_fitness(kg)
  137. if fg > f:
  138. f, k = fg, kg.copy()
  139. pbar.desc = f'{PREFIX}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
  140. if verbose:
  141. print_results(k, verbose)
  142. return print_results(k)