import torch import torch.nn as nn import math from ultralytics import YOLO from ultralytics.nn.modules import Conv, Concat from lib.models.common import Focus, BottleneckCSP, Detect from lib.utils import check_anchor_order import logging class YOLOv11Backbone(nn.Module): def __init__(self, width_multiple=0.25, depth_multiple=0.50, yolo_model_path=None): """ YOLOv11 Backbone - 直接从 ultralytics YOLO 模型提取 Args: width_multiple: 通道数缩放因子 (n=0.25, s=0.50, m=1.00, l=1.00, x=1.50) depth_multiple: 深度缩放因子 (n=0.50, s=0.50, m=0.50, l=1.00, x=1.00) yolo_model_path: YOLOv11 预训练模型路径（可选） Warning: 不同的yolo model(n, s, m, l, x)模型结构都会不同，目前这个是以 small 为例， 恰好可以输出(128, 256, 512)通道数 (虽然有adapter也无所谓) """ super().__init__() self.out_indices = [4, 6, 10] # P3, P4, P5 # 如果提供了预训练模型路径，直接加载 if yolo_model_path: yolo = YOLO(yolo_model_path) yolo_model = yolo.model # 提取 backbone 层 (0-10) self.layers = nn.ModuleList([yolo_model.model[i] for i in range(11)]) # 获取输出通道数 (C3k2 和 C2PSA 都有 cv2 属性) self.out_channels = [ yolo_model.model[self.out_indices[0]].conv.out_channels, # P3 (C3k2) yolo_model.model[self.out_indices[1]].conv.out_channels, # P4 (C3k2) yolo_model.model[self.out_indices[2]].conv.out_channels, # P5 (C2PSA) ] else: # 如果没有预训练模型，使用 ultralytics 的模块构建 from ultralytics.nn.modules import Conv, C3k2, SPPF, C2PSA # 根据 width_multiple 计算通道数 def make_divisible(x, divisor=8): """确保通道数是 divisor 的倍数""" return int(math.ceil(x / divisor) * divisor) c1 = make_divisible(64 * width_multiple) c2 = make_divisible(128 * width_multiple) c3 = make_divisible(256 * width_multiple) c4 = make_divisible(512 * width_multiple) c5 = make_divisible(1024 * width_multiple) # 根据 depth_multiple 计算重复次数 n1 = max(round(2 * depth_multiple), 1) # C3k2 repeats self.layers = nn.ModuleList([ Conv(3, c1, k=3, s=2), # 0 Conv(c1, c2, k=3, s=2), # 1 C3k2(c2, c3, n=n1, shortcut=False, e=0.25), # 2 Conv(c3, c3, k=3, s=2), # 3 C3k2(c3, c4, n=n1, shortcut=False, e=0.25), # 4 Conv(c4, c4, k=3, s=2), # 5 C3k2(c4, c4, n=n1, shortcut=True), # 6 Conv(c4, c5, k=3, s=2), # 7 C3k2(c5, c5, n=n1, shortcut=True), # 8 SPPF(c5, c5, k=5), # 9 C2PSA(c5, c5, n=n1), # 10 ]) self.out_channels = [] for i in self.out_indices: layer = self.layers[i] #（Conv) if hasattr(layer, 'conv'): self.out_channels.append(layer.conv.out_channels) elif hasattr(layer, 'cv2'): # (C3k2) self.out_channels.append(layer.cv2.conv.out_channels) else: raise AttributeError(f"Layer {i} 没有 conv 或 cv2 属性，请检查模块结构") def forward(self, x): outputs = [] for i, layer in enumerate(self.layers): x = layer(x) if i in self.out_indices: outputs.append(x) return outputs class ChannelAdapter(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return self.conv(x) class YOLOPWithYOLOv11(nn.Module): def __init__(self, num_seg_class=2, yolo_scale='n', yolo_weights_path=None): """ YOLOP with YOLOv11 Backbone Args: num_seg_class: 分割类别数 yolo_scale: YOLOv11 规模 ('n', 's', 'm', 'l', 'x') yolo_weights_path: YOLOv11 预训练权重路径（可选） """ super().__init__() # YOLOv11 缩放参数 scale_configs = { 'n': {'width': 0.25, 'depth': 0.50}, # nano 's': {'width': 0.50, 'depth': 0.50}, # small 'm': {'width': 1.00, 'depth': 0.50}, # medium 'l': {'width': 1.00, 'depth': 1.00}, # large 'x': {'width': 1.50, 'depth': 1.00}, # xlarge } if yolo_scale not in scale_configs: raise ValueError(f"Invalid yolo_scale: {yolo_scale}. Must be one of {list(scale_configs.keys())}") scale = scale_configs[yolo_scale] # 如果提供了权重路径，直接从预训练模型提取 backbone if yolo_weights_path: self.backbone = YOLOv11Backbone(yolo_model_path=yolo_weights_path) else: self.backbone = YOLOv11Backbone(width_multiple=scale['width'], depth_multiple=scale['depth']) # 适配 YOLOv11 输出到 YOLOP neck 输入 [128, 256, 512] backbone_channels = self.backbone.out_channels neck_channels = [128, 256, 512] self.adapters = nn.ModuleList([ ChannelAdapter(backbone_channels[0], neck_channels[0]), # P3 ChannelAdapter(backbone_channels[1], neck_channels[1]), # P4 ChannelAdapter(backbone_channels[2], neck_channels[2]), # P5 ]) # YOLOP neck (层 11-24) self.neck = nn.ModuleList([ Conv(512, 256, k=1, s=1), # 11 nn.Upsample(scale_factor=2, mode='nearest'), # 12 Concat(dimension=1), # 13: Concat [-1, 6] BottleneckCSP(512, 256, n=1, shortcut=False), # 14 Conv(256, 128, k=1, s=1), # 15 nn.Upsample(scale_factor=2, mode='nearest'), # 16 Concat(dimension=1), # 17: Concat [-1, 4] BottleneckCSP(256, 128, n=1, shortcut=False), # 18 Conv(128, 128, k=3, s=2), # 19 Concat(dimension=1), # 20: Concat [-1, 14] BottleneckCSP(256, 256, n=1, shortcut=False), # 21 Conv(256, 256, k=3, s=2), # 22 Concat(dimension=1), # 23: Concat [-1, 10] BottleneckCSP(512, 512, n=1, shortcut=False), # 24 ]) # YOLOP heads self.detect_head = Detect(1, [[3,9,5,11,4,20], [7,18,6,39,12,31], [19,50,38,81,68,157]], [128, 256, 512]) self.drivable_seg_head = nn.ModuleList([ Conv(256, 128, k=3, s=1), # 25 nn.Upsample(scale_factor=2, mode='nearest'), # 26 BottleneckCSP(128, 64, n=1, shortcut=False), # 27 Conv(64, 32, k=3, s=1), # 28 nn.Upsample(scale_factor=2, mode='nearest'), # 29 Conv(32, 16, k=3, s=1), # 30 BottleneckCSP(16, 8, n=1, shortcut=False), # 31 nn.Upsample(scale_factor=2, mode='nearest'), # 32 Conv(8, num_seg_class, k=3, s=1), # 33 ]) self.lane_seg_head = nn.ModuleList([ Conv(256, 128, k=3, s=1), # 34 nn.Upsample(scale_factor=2, mode='nearest'), # 35 BottleneckCSP(128, 64, n=1, shortcut=False), # 36 Conv(64, 32, k=3, s=1), # 37 nn.Upsample(scale_factor=2, mode='nearest'), # 38 Conv(32, 16, k=3, s=1), # 39 BottleneckCSP(16, 8, n=1, shortcut=False), # 40 nn.Upsample(scale_factor=2, mode='nearest'), # 41 Conv(8, 2, k=3, s=1), # 42 ]) # 初始化 Detection Head 的 stride # self.detect_head.stride = torch.tensor([8., 16., 32.]) # P3, P4, P5 的 stride # 初始化时动态计算 stride s = 128 with torch.no_grad(): dummy = torch.zeros(1, 3, s, s) detect_out, _, _ = self.forward(dummy) self.detect_head.stride = torch.tensor([s / x.shape[-2] for x in detect_out]) self.detect_head.anchors /= self.detect_head.stride.view(-1, 1, 1) # Set the anchors for the corresponding scale check_anchor_order(self.detect_head) self.stride = self.detect_head.stride print(f"Initialized Detect head with strides: {self.detect_head.stride.tolist()}") # 添加必要的属性以兼容训练代码 self.nc = 1 # number of classes self.detector_index = -1 # detector在模型中的索引 self.names = ['vehicle'] # class names self.model = nn.ModuleList([ self.backbone, self.adapters, self.neck, self.detect_head, self.drivable_seg_head, self.lane_seg_head ]) self.detector_index = 3 # detect_head 在第4个位置 self.det_out_idx = 25 self.gr = 1.0 # giou loss ratio (obj loss ratio is 1-giou) # 初始化 Detection Head 的偏置 self._initialize_biases() def freeze_backbone(self): """冻结backbone和adapters的参数""" logging.info("Freezing backbone parameters...") for param in self.backbone.parameters(): param.requires_grad = False for param in self.adapters.parameters(): param.requires_grad = False # 验证冻结状态 frozen_count = sum(1 for p in self.backbone.parameters() if not p.requires_grad) frozen_count += sum(1 for p in self.adapters.parameters() if not p.requires_grad) total_count = sum(1 for _ in self.backbone.parameters()) total_count += sum(1 for _ in self.adapters.parameters()) logging.info(f"Frozen {frozen_count}/{total_count} backbone+adapter parameters") def unfreeze_backbone(self): """解冻backbone和adapters的参数""" logging.info("Unfreezing backbone parameters...") for param in self.backbone.parameters(): param.requires_grad = True for param in self.adapters.parameters(): param.requires_grad = True def _initialize_biases(self, cf=None): """初始化检测头的偏置 (参考原始YOLOP实现)""" # https://arxiv.org/abs/1708.02002 section 3.3 m = self.detect_head # Detect() module for mi, s in zip(m.m, m.stride): # from b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) def load_yolov11_backbone_weights(self, weights_path, freeze_backbone=False): """ 从YOLOv11预训练模型加载backbone权重 Args: weights_path: YOLOv11权重路径（.pt文件） freeze_backbone: 是否冻结backbone参数 """ try: from ultralytics import YOLO logging.info(f"Loading YOLOv11 weights from {weights_path}") # 加载YOLOv11模型 yolo_model = YOLO(weights_path) yolo_state_dict = yolo_model.model.state_dict() # 映射YOLOv11的backbone权重到我们的模型 # YOLOv11的backbone层索引: 0-10 backbone_mapping = { # YOLOv11 layer -> our layer 'model.0': 'backbone.layers.0', # Conv 3->64 'model.1': 'backbone.layers.1', # Conv 64->128 'model.2': 'backbone.layers.2', # C3k2 128->256 'model.3': 'backbone.layers.3', # Conv 256->256 'model.4': 'backbone.layers.4', # C3k2 256->512 'model.5': 'backbone.layers.5', # Conv 512->512 'model.6': 'backbone.layers.6', # C3k2 512->512 'model.7': 'backbone.layers.7', # Conv 512->1024 'model.8': 'backbone.layers.8', # C3k2 1024->1024 'model.9': 'backbone.layers.9', # SPPF 'model.10': 'backbone.layers.10', # C2PSA } # 构建新的state dict new_state_dict = {} loaded_keys = [] for yolo_key, our_key in backbone_mapping.items(): for k, v in yolo_state_dict.items(): if k.startswith(yolo_key + '.'): new_key = k.replace(yolo_key, our_key) new_state_dict[new_key] = v loaded_keys.append(new_key) # 加载权重 model_dict = self.state_dict() # 只更新存在的键 new_state_dict = {k: v for k, v in new_state_dict.items() if k in model_dict} model_dict.update(new_state_dict) self.load_state_dict(model_dict) logging.info(f"Successfully loaded {len(loaded_keys)} backbone parameters from YOLOv11") # 冻结backbone if freeze_backbone: self.freeze_backbone() logging.info("Backbone frozen successfully") except Exception as e: logging.warning(f"Failed to load YOLOv11 weights: {e}") logging.warning("Training will start from scratch") def forward(self, x): features = self.backbone(x) # YOLOv11 输出 [P3, P4, P5] features = [adapter(f) for adapter, f in zip(self.adapters, features)] # 适配到 [128, 256, 512] # Neck 前向传播 x = features[-1] # P5 10 x = self.neck[0](x) # 11 x = self.neck[1](x) # 12 x = self.neck[2]([x, features[1]]) # 13 x = self.neck[3](x) # 14 x = self.neck[4](x) # 15 x = self.neck[5](x) # 16 p3_fpn = self.neck[6]([x, features[0]]) # 17 (P3, 256 通道) p3 = self.neck[7](p3_fpn) # 18 (P3, 128 通道) x = self.neck[8](p3) # 19 x = self.neck[9]([x, self.neck[4](features[1])]) # 20 p4 = self.neck[10](x) # 21 x = self.neck[11](p4) # 22 x = self.neck[12]([x, self.neck[0](features[2])]) # 23 p5 = self.neck[13](x) # 24 # Heads detect_out = self.detect_head([p3, p4, p5]) # 使用层 17, 20, 23 drivable_out = p3_fpn # 使用层 16 for layer in self.drivable_seg_head: drivable_out = layer(drivable_out) lane_out = p3_fpn # 使用层 16 for layer in self.lane_seg_head: lane_out = layer(lane_out) drivable_out = torch.sigmoid(drivable_out) lane_out = torch.sigmoid(lane_out) return [detect_out, drivable_out, lane_out] def get_net_yolov11(cfg, **kwargs): """ 获取带有YOLOv11 backbone的YOLOP模型 Args: cfg: 配置对象 **kwargs: 其他参数，包括： - yolov11_weights: YOLOv11预训练权重路径 - freeze_backbone: 是否冻结backbone - yolo_scale: YOLOv11规模 ('n', 's', 'm', 'l', 'x') """ num_seg_class = cfg.num_seg_class if hasattr(cfg, 'num_seg_class') else 2 yolo_scale = kwargs.get('yolo_scale', 'n') # 默认使用 nano # 如果提供了权重路径，直接用权重初始化 yolov11_weights = kwargs.get('yolov11_weights', f'weights/yolo11{yolo_scale}.pt') freeze_backbone = kwargs.get('freeze_backbone', False) # 在初始化时就加载预训练权重 import os if os.path.exists(yolov11_weights): logging.info(f"Creating model with YOLOv11{yolo_scale} pretrained weights from {yolov11_weights}") model = YOLOPWithYOLOv11(num_seg_class=num_seg_class, yolo_scale=yolo_scale, yolo_weights_path=yolov11_weights) if freeze_backbone: model.freeze_backbone() else: logging.warning(f"YOLOv11 weights not found at {yolov11_weights}, creating model from scratch") model = YOLOPWithYOLOv11(num_seg_class=num_seg_class, yolo_scale=yolo_scale, yolo_weights_path=None) return model