modify init

yolo/models/yolov5_af/yolov5_af_basic.py

@@ -63,9 +63,9 @@ class BasicConv(nn.Module):
             return self.act(self.norm(self.conv(x)))
         else:
             # Depthwise conv
-            x = self.norm1(self.conv1(x))
+            x = self.act(self.norm1(self.conv1(x)))
             # Pointwise conv
-            x = self.norm2(self.conv2(x))
+            x = self.act(self.norm2(self.conv2(x)))
             return x
 
 

yolo/models/yolov5_af/yolov5_af_neck.py

@@ -28,6 +28,17 @@ class SPPF(nn.Module):
                               stride=1,
                               padding=cfg.spp_pooling_size // 2)
 
+        # Initialize all layers
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, x):
         x = self.cv1(x)
         y1 = self.m(x)

yolo/models/yolov5_af/yolov5_af_pafpn.py

@@ -78,6 +78,17 @@ class Yolov5PaFPN(nn.Module):
                       ])
         self.out_dims = [round(cfg.head_dim*cfg.width)] * 3
 
+        # Initialize all layers
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, features):
         c3, c4, c5 = features
         

yolo/models/yolov7_af/yolov7_af_neck.py

@@ -24,6 +24,17 @@ class SPPF(nn.Module):
                               stride=1,
                               padding=cfg.spp_pooling_size // 2)
 
+        # Initialize all layers
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, x):
         x = self.cv1(x)
         y1 = self.m(x)
@@ -51,7 +62,17 @@ class SPPFBlockCSP(nn.Module):
                       )
         self.cv3 = BasicConv(inter_dim * 2, self.out_dim, kernel_size=1, act_type=cfg.neck_act, norm_type=cfg.neck_norm)
 
-        
+        # Initialize all layers
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, x):
         x1 = self.cv1(x)
         x2 = self.module(self.cv2(x))

yolo/models/yolov7_af/yolov7_af_pafpn.py

@@ -85,6 +85,17 @@ class Yolov7PaFPN(nn.Module):
                                      kernel_size=3, padding=1, stride=1,
                                      act_type=cfg.fpn_act, norm_type=cfg.fpn_norm, depthwise=cfg.fpn_depthwise)
 
+        # Initialize all layers
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, features):
         c3, c4, c5 = features
 

yolo/models/yolov8/yolov8_neck.py

@@ -28,6 +28,17 @@ class SPPF(nn.Module):
                               stride=1,
                               padding=cfg.spp_pooling_size // 2)
 
+        # Initialize all layers
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, x):
         x = self.cv1(x)
         y1 = self.m(x)