diff --git a/models/common.py b/models/common.py index 341e878..0dec4be 100644 --- a/models/common.py +++ b/models/common.py @@ -1121,3 +1121,25 @@ 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 ECA(nn.Module): + """Efficient Channel Attention""" + def __init__(self, channels, k_size=3): + super().__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.conv = nn.Conv1d( + 1, 1, kernel_size=k_size, + padding=(k_size - 1) // 2, + bias=False + ) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + # x: [B, C, H, W] + y = self.avg_pool(x) # [B, C, 1, 1] + y = y.squeeze(-1).transpose(-1, -2) # [B, 1, C] + y = self.conv(y) # [B, 1, C] + y = self.sigmoid(y) + y = y.transpose(-1, -2).unsqueeze(-1) # [B, C, 1, 1] + return x * y diff --git a/models/yolov5s.yaml b/models/yolov5s.yaml index 11ff790..17689e0 100644 --- a/models/yolov5s.yaml +++ b/models/yolov5s.yaml @@ -25,25 +25,54 @@ backbone: [-1, 1, SPPF, [1024, 5]], # 9 ] -# YOLOv5 v6.0 head -head: [ - [-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]], # 13 +# # YOLOv5 v6.0 head +# head: [ +# [-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]], # 13 - [-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]], # 17 (P3/8-small) +# [-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]], # 17 (P3/8-small) - [-1, 1, Conv, [256, 3, 2]], - [[-1, 14], 1, Concat, [1]], # cat head P4 - [-1, 3, C3, [512, False]], # 20 (P4/16-medium) +# [-1, 1, Conv, [256, 3, 2]], +# [[-1, 14], 1, Concat, [1]], # cat head P4 +# [-1, 3, C3, [512, False]], # 20 (P4/16-medium) - [-1, 1, Conv, [512, 3, 2]], - [[-1, 10], 1, Concat, [1]], # cat head P5 - [-1, 3, C3, [1024, False]], # 23 (P5/32-large) +# [-1, 1, Conv, [512, 3, 2]], +# [[-1, 10], 1, Concat, [1]], # cat head P5 +# [-1, 3, C3, [1024, False]], # 23 (P5/32-large) - [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) +# [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) +# ] + +head: + [ + # -------- Top-down pathway (FPN) -------- + [-1, 1, Conv, [512, 1, 1]], # 10 reduce channels (P5) + [-1, 1, nn.Upsample, [None, 2, "nearest"]], # 11 upsample + [[-1, 6], 1, Concat, [1]], # 12 concat backbone P4 + [-1, 3, C3, [512, False]], # 13 P4 fused feature + + [-1, 1, Conv, [256, 1, 1]], # 14 reduce channels + [-1, 1, nn.Upsample, [None, 2, "nearest"]], # 15 upsample + [[-1, 4], 1, Concat, [1]], # 16 concat backbone P3 + [-1, 3, C3, [256, False]], # 17 P3 fused feature + [-1, 1, ECA, [256]], # 18 ECA on P3 (semantic calibration) + + # -------- Bottom-up pathway (PAN) -------- + [-1, 1, Conv, [256, 3, 2]], # 19 downsample + [[-1, 14], 1, Concat, [1]], # 20 concat head P4 + [-1, 3, C3, [512, False]], # 21 P4 fused feature + [-1, 1, ECA, [512]], # 22 ECA on P4 + + [-1, 1, Conv, [512, 3, 2]], # 23 downsample + [[-1, 10], 1, Concat, [1]], # 24 concat head P5 + [-1, 3, C3, [1024, False]], # 25 P5 fused feature + [-1, 1, ECA, [1024]], # 26 ECA on P5 + + # -------- Detection heads -------- + [[18, 22, 26], 1, Detect, [nc, anchors]], # 27 Detect(P3, P4, P5) ] diff --git a/train_my_gra.sh b/train_my_gra.sh new file mode 100644 index 0000000..37f6dc6 --- /dev/null +++ b/train_my_gra.sh @@ -0,0 +1,9 @@ +export CUDA_VISIBLE_DEVICES=4,5,6,7 +python /data6/zhangw/yolov5/train.py \ + --img 960 \ + --batch 16 \ + --epochs 300 \ + --data /data6/zhangw/ZW_gra_project/CUB_100/cub100.yaml \ + --weights /data6/zhangw/yolov5/yolov5s6.pt \ + --project /data6/zhangw/ZW_gra_project/yolov5_result \ + --name G2_cub100_yolov5s6_changeNeck_e300 \ No newline at end of file