diff --git a/models/common.py b/models/common.py index 341e878..4df5a6a 100644 --- a/models/common.py +++ b/models/common.py @@ -1121,3 +1121,36 @@ class Classify(nn.Module): if isinstance(x, list): x = torch.cat(x, 1) return self.linear(self.drop(self.pool(self.conv(x)).flatten(1))) + + +# 新增 +class FGEMSafe(nn.Module): + """Safer lightweight fine-grained enhancement block for cls branch.""" + def __init__(self, c1, c_mid=None, residual_scale=0.1): + super().__init__() + c_mid = c_mid or max(c1 // 2, 32) + self.residual_scale = residual_scale + + # 1x1 reduce + self.reduce = Conv(c1, c_mid, 1, 1) + + # depthwise 3x3 local modeling + self.dw = nn.Sequential( + nn.Conv2d(c_mid, c_mid, kernel_size=3, stride=1, padding=1, groups=c_mid, bias=False), + nn.BatchNorm2d(c_mid), + nn.SiLU() + ) + + # pointwise fusion + self.pw = Conv(c_mid, c_mid, 1, 1) + + # 1x1 expand + self.expand = Conv(c_mid, c1, 1, 1) + + def forward(self, x): + identity = x + y = self.reduce(x) + y = self.dw(y) + y = self.pw(y) + y = self.expand(y) + return identity + self.residual_scale * y \ No newline at end of file diff --git a/models/yolo.py b/models/yolo.py index 1198e9c..ee985a0 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -48,6 +48,7 @@ from models.common import ( GhostBottleneck, GhostConv, Proto, + FGEMSafe, ) from models.experimental import MixConv2d from utils.autoanchor import check_anchor_order @@ -152,7 +153,23 @@ class DetectDecoupledLite(nn.Module): self.regobj_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch) # cls trunk - self.cls_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch) + # self.cls_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch) + # cls trunk: only enhance higher-level scales for large-object bird recognition + self.cls_convs = nn.ModuleList() + for i, _ in enumerate(ch): + if i < 2: + # P3, P4: keep original cls trunk + self.cls_convs.append( + Conv(head_c, head_c, 3, 1) + ) + else: + # P5, P6: add safe FG enhancement + self.cls_convs.append( + nn.Sequential( + Conv(head_c, head_c, 3, 1), + FGEMSafe(head_c, head_c // 2, residual_scale=0.1) + ) + ) # predictors self.box_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * 4, 1) for _ in ch) diff --git a/models/yolov5s6_decoupled_light_FGEM_safe.yaml b/models/yolov5s6_decoupled_light_FGEM_safe.yaml new file mode 100644 index 0000000..39a255c --- /dev/null +++ b/models/yolov5s6_decoupled_light_FGEM_safe.yaml @@ -0,0 +1,62 @@ +# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license + +# Parameters +nc: 100 # number of classes +depth_multiple: 0.33 # model depth multiple +width_multiple: 0.50 # layer channel multiple +anchors: + - [19, 27, 44, 40, 38, 94] # P3/8 + - [96, 68, 86, 152, 180, 137] # P4/16 + - [140, 301, 303, 264, 238, 542] # P5/32 + - [436, 615, 739, 380, 925, 792] # P6/64 + +# YOLOv5 v6.0 backbone +backbone: + # [from, number, module, args] + [ + [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 + [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 + [-1, 3, C3, [128]], + [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 + [-1, 6, C3, [256]], + [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 + [-1, 9, C3, [512]], + [-1, 1, Conv, [768, 3, 2]], # 7-P5/32 + [-1, 3, C3, [768]], + [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64 + [-1, 3, C3, [1024]], + [-1, 1, SPPF, [1024, 5]], # 11 + ] + +# YOLOv5 v6.0 head +head: [ + [-1, 1, Conv, [768, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, "nearest"]], + [[-1, 8], 1, Concat, [1]], # cat backbone P5 + [-1, 3, C3, [768, False]], # 15 + + [-1, 1, Conv, [512, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, "nearest"]], + [[-1, 6], 1, Concat, [1]], # cat backbone P4 + [-1, 3, C3, [512, False]], # 19 + + [-1, 1, Conv, [256, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, "nearest"]], + [[-1, 4], 1, Concat, [1]], # cat backbone P3 + [-1, 3, C3, [256, False]], # 23 (P3/8-small) + + [-1, 1, Conv, [256, 3, 2]], + [[-1, 20], 1, Concat, [1]], # cat head P4 + [-1, 3, C3, [512, False]], # 26 (P4/16-medium) + + [-1, 1, Conv, [512, 3, 2]], + [[-1, 16], 1, Concat, [1]], # cat head P5 + [-1, 3, C3, [768, False]], # 29 (P5/32-large) + + [-1, 1, Conv, [768, 3, 2]], + [[-1, 12], 1, Concat, [1]], # cat head P6 + [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge) + + # [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6) + [[23, 26, 29, 32], 1, DetectDecoupledLite, [nc, anchors]], # Detect(P3, P4, P5, P6) + ] \ No newline at end of file diff --git a/train_my_gra.sh b/train_my_gra.sh index e06572c..176acd7 100644 --- a/train_my_gra.sh +++ b/train_my_gra.sh @@ -6,5 +6,5 @@ python /data6/zhangw/yolov5/train.py \ --data /data6/zhangw/ZW_gra_project/CUB_100/cub100.yaml \ --weights /data6/zhangw/yolov5/yolov5s6.pt \ --project /data6/zhangw/ZW_gra_project/yolov5_result_pretrained \ - --name G8_decoupled_head_light_e300 \ - --cfg /data6/zhangw/yolov5/models/yolov5s6_decoupled_light.yaml \ No newline at end of file + --name G10_decoupled_head_light_FGEM_safe_e300 \ + --cfg /data6/zhangw/yolov5/models/yolov5s6_decoupled_light_FGEM_safe.yaml \ No newline at end of file diff --git a/train_my_gra_resume.sh b/train_my_gra_resume.sh new file mode 100644 index 0000000..8372600 --- /dev/null +++ b/train_my_gra_resume.sh @@ -0,0 +1,6 @@ +export CUDA_VISIBLE_DEVICES=4,5,6,7 +python /data6/zhangw/yolov5/train.py \ + --resume /data6/zhangw/ZW_gra_project/yolov5_result_pretrained/G10_decoupled_head_light_FGEM_safe_e1/weights/last.pt \ + --epochs 2 \ + --project /data6/zhangw/ZW_gra_project/yolov5_result_pretrained \ + --name G10_decoupled_head_light_FGEM_safe_e2_resume \ No newline at end of file