commit

yolov6/models/yolo.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+import math
+import torch.nn as nn
+from yolov6.layers.common import *
+from yolov6.utils.torch_utils import initialize_weights
+from yolov6.models.efficientrep import EfficientRep
+from yolov6.models.reppan import RepPANNeck
+from yolov6.models.effidehead import Detect, build_effidehead_layer
+
+
+class Model(nn.Module):
+    '''YOLOv6 model with backbone, neck and head.
+    The default parts are EfficientRep Backbone, Rep-PAN and
+    Efficient Decoupled Head.
+    '''
+    def __init__(self, config, channels=3, num_classes=None, anchors=None):  # model, input channels, number of classes
+        super().__init__()
+        # Build network
+        num_layers = config.model.head.num_layers
+        self.backbone, self.neck, self.detect = build_network(config, channels, num_classes, anchors, num_layers)
+
+        # Init Detect head
+        begin_indices = config.model.head.begin_indices
+        out_indices_head = config.model.head.out_indices
+        self.stride = self.detect.stride
+        self.detect.i = begin_indices
+        self.detect.f = out_indices_head
+        self.detect.initialize_biases()
+
+        # Init weights
+        initialize_weights(self)
+
+    def forward(self, x):
+        x = self.backbone(x)
+        x = self.neck(x)
+        x = self.detect(x)
+        return x
+
+    def _apply(self, fn):
+        self = super()._apply(fn)
+        self.detect.stride = fn(self.detect.stride)
+        self.detect.grid = list(map(fn, self.detect.grid))
+        return self
+
+
+def make_divisible(x, divisor):
+    # Upward revision the value x to make it evenly divisible by the divisor.
+    return math.ceil(x / divisor) * divisor
+
+
+def build_network(config, channels, num_classes, anchors, num_layers):
+    depth_mul = config.model.depth_multiple
+    width_mul = config.model.width_multiple
+    num_repeat_backbone = config.model.backbone.num_repeats
+    channels_list_backbone = config.model.backbone.out_channels
+    num_repeat_neck = config.model.neck.num_repeats
+    channels_list_neck = config.model.neck.out_channels
+    num_anchors = config.model.head.anchors
+    num_repeat = [(max(round(i * depth_mul), 1) if i > 1 else i) for i in (num_repeat_backbone + num_repeat_neck)]
+    channels_list = [make_divisible(i * width_mul, 8) for i in (channels_list_backbone + channels_list_neck)]
+
+    backbone = EfficientRep(
+        in_channels=channels,
+        channels_list=channels_list,
+        num_repeats=num_repeat
+    )
+
+    neck = RepPANNeck(
+        channels_list=channels_list,
+        num_repeats=num_repeat
+    )
+
+    head_layers = build_effidehead_layer(channels_list, num_anchors, num_classes)
+
+    head = Detect(num_classes, anchors, num_layers, head_layers=head_layers)
+
+    return backbone, neck, head
+
+
+def build_model(cfg, num_classes, device):
+    model = Model(cfg, channels=3, num_classes=num_classes, anchors=cfg.model.head.anchors).to(device)
+    return model