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.
84 lines
2.9 KiB
84 lines
2.9 KiB
#!/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
|