[Feature] Add autoanchor for YOLO series

configs/yolov5/autoanchor/yolov5_s-v61_syncbn_8xb16-300e_coco_aa.py

@@ -0,0 +1,13 @@
+_base_ = '../yolov5_s-v61_syncbn_8xb16-300e_coco.py'
+
+model = dict(
+    bbox_head=dict(prior_generator=dict(type='YOLOAutoAnchorGenerator')))
+
+custom_hooks = [
+    dict(
+        type='YOLOAutoAnchorHook',
+        optimizer=dict(
+            type='YOLOKMeansAnchorOptimizer',
+            iters=1000,
+            num_anchor_per_level=[3, 3, 3]))
+]

mmyolo/core/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .anchor import *  # noqa: F401, F403

mmyolo/core/anchor/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .anchor_generator import YOLOAutoAnchorGenerator
+from .anchor_optimizer import (YOLODEAnchorOptimizer,
+                               YOLOKMeansAnchorOptimizer,
+                               YOLOV5KMeansAnchorOptimizer)
+
+__all__ = [
+    'YOLOAutoAnchorGenerator',
+    'YOLOKMeansAnchorOptimizer',
+    'YOLOV5KMeansAnchorOptimizer',
+    'YOLODEAnchorOptimizer',
+]

mmyolo/core/anchor/anchor_generator.py

@@ -0,0 +1,40 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+import torch.nn as nn
+from mmdet.models.task_modules import YOLOAnchorGenerator
+from torch.nn.modules.utils import _pair
+
+from mmyolo.registry import TASK_UTILS
+
+
+@TASK_UTILS.register_module()
+class YOLOAutoAnchorGenerator(nn.Module, YOLOAnchorGenerator):
+    """AutoAnchor generator for YOLO.
+
+    Args:
+        strides (list[int] | list[tuple[int, int]]): Strides of anchors
+            in multiple feature levels.
+        base_sizes (list[list[tuple[int, int]]]): The basic sizes
+            of anchors in multiple levels.
+    """
+
+    def __init__(self, strides, base_sizes, use_box_type: bool = False):
+        super().__init__()
+        self.strides = [_pair(stride) for stride in strides]
+        self.centers = [(stride[0] / 2., stride[1] / 2.)
+                        for stride in self.strides]
+        self.use_box_type = use_box_type
+        self.register_buffer('anchors', torch.tensor(base_sizes))
+
+    @property
+    def base_sizes(self):
+        T = []
+        num_anchor_per_level = len(self.anchors[0])
+        for base_sizes_per_level in self.anchors:
+            assert num_anchor_per_level == len(base_sizes_per_level)
+            T.append([_pair(base_size) for base_size in base_sizes_per_level])
+        return T
+
+    @property
+    def base_anchors(self):
+        return self.gen_base_anchors()