From 82415035801dd586810495b4e6970db37a1e1da7 Mon Sep 17 00:00:00 2001
From: Zhang Wei <3160665918@qq.com>
Date: Wed, 24 Dec 2025 16:29:39 +0800
Subject: [PATCH] <Neck_2>change Concat to WeightedAdd

---
 models/common.py    | 20 +++++++++++++++
 models/yolo.py      |  6 +++++
 models/yolov5s.yaml | 59 +++++++++++++++++++++++++++++++--------------
 train_my_gra.sh     |  4 +--
 4 files changed, 69 insertions(+), 20 deletions(-)

diff --git a/models/common.py b/models/common.py
index 341e878..99bb95e 100644
--- a/models/common.py
+++ b/models/common.py
@@ -1121,3 +1121,23 @@ 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 WeightedAdd(nn.Module):
+    """
+    Weighted feature fusion (BiFPN-style)
+    inputs: list of feature maps with same shape
+    """
+    def __init__(self, n_inputs, eps=1e-4):
+        super().__init__()
+        self.eps = eps
+        self.w = nn.Parameter(torch.ones(n_inputs), requires_grad=True)
+
+    def forward(self, x):
+        # x: list of tensors
+        w = torch.relu(self.w)
+        w = w / (w.sum() + self.eps)
+        out = 0
+        for i in range(len(x)):
+            out += w[i] * x[i]
+        return out
diff --git a/models/yolo.py b/models/yolo.py
index 13498ac..3a59474 100644
--- a/models/yolo.py
+++ b/models/yolo.py
@@ -37,6 +37,7 @@ from models.common import (
     C3x,
     Classify,
     Concat,
+    WeightedAdd,   # ✅ 新增
     Contract,
     Conv,
     CrossConv,
@@ -433,6 +434,11 @@ def parse_model(d, ch):
             args = [ch[f]]
         elif m is Concat:
             c2 = sum(ch[x] for x in f)
+            
+        elif m is WeightedAdd:
+            # 加权相加：所有输入通道必须一致
+            c2 = ch[f[0]]
+
         # TODO: channel, gw, gd
         elif m in {Detect, Segment}:
             args.append([ch[x] for x in f])
diff --git a/models/yolov5s.yaml b/models/yolov5s.yaml
index 11ff790..9082635 100644
--- a/models/yolov5s.yaml
+++ b/models/yolov5s.yaml
@@ -25,25 +25,48 @@ 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
+  [-1, 1, Conv, [512, 1, 1]]
+  [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  [[-1, 6], 1, WeightedAdd, [2]]
+  [-1, 3, C3, [512, False]]
+
+  [-1, 1, Conv, [256, 1, 1]]
+  [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  [[-1, 4], 1, WeightedAdd, [2]]
+  [-1, 3, C3, [256, False]]
+
+  # Bottom-up
+  [-1, 1, Conv, [256, 3, 2]]
+  [[-1, 14], 1, WeightedAdd, [2]]
+  [-1, 3, C3, [512, False]]
+
+  [-1, 1, Conv, [512, 3, 2]]
+  [[-1, 10], 1, WeightedAdd, [2]]
+  [-1, 3, C3, [1024, False]]
+
+  [[17, 20, 23], 1, Detect, [nc, anchors]]
\ No newline at end of file
diff --git a/train_my_gra.sh b/train_my_gra.sh
index f20bf38..5c036bb 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 \
-    --name G0_cub100_yolov5s6_raw_e300 \
-    --cfg /data6/zhangw/yolov5/models/yolov5s.yaml
\ No newline at end of file
+    --cfg /data6/zhangw/yolov5/models/yolov5s.yaml \
+    --name G2_2_cub100_yolov5s6_changeNeck_e300
\ No newline at end of file