modify init

yolo/dataset/build.py

@@ -25,7 +25,6 @@ def build_dataset(args, cfg, transform=None, is_train=False):
     ## VOC dataset
     if args.dataset == 'voc':
         image_set = [('2007', 'trainval'), ('2012', 'trainval')] if is_train else [('2007', 'test')]
-        cfg.num_classes  = 20
         dataset = VOCDataset(cfg       = cfg,
                              data_dir  = args.root,
                              image_set = image_set,
@@ -35,7 +34,6 @@ def build_dataset(args, cfg, transform=None, is_train=False):
     ## COCO dataset
     elif args.dataset == 'coco':
         image_set = 'train2017' if is_train else 'val2017'
-        cfg.num_classes  = 80
         dataset = COCODataset(cfg       = cfg,
                               data_dir  = args.root,
                               image_set = image_set,
@@ -45,7 +43,6 @@ def build_dataset(args, cfg, transform=None, is_train=False):
     ## Custom dataset
     elif args.dataset == 'customed':
         image_set = 'train' if is_train else 'val'
-        cfg.num_classes  = 20
         dataset = CustomedDataset(cfg       = cfg,
                                   data_dir  = args.root,
                                   image_set = image_set,

yolo/dataset/customed.py

@@ -28,6 +28,7 @@ class CustomedDataset(Dataset):
         self.image_set = image_set
         self.is_train  = is_train
         self.num_classes = len(customed_class_labels)
+        self.num_classes = 9
         # ----------- Path parameters -----------
         self.data_dir = data_dir
         self.json_file = '{}.json'.format(image_set)

yolo/dataset/voc.py

@@ -54,7 +54,7 @@ class VOCDataset(data.Dataset):
         # ----------- Basic parameters -----------
         self.image_set = image_set
         self.is_train  = is_train
-        self.num_classes = 80
+        self.num_classes = 20
         # ----------- Path parameters -----------
         self.root = data_dir
         self._annopath = osp.join('%s', 'Annotations', '%s.xml')

yolo/engine.py

@@ -275,7 +275,7 @@ class YoloTrainer(object):
                 # refine tgt
                 tgt_boxes_wh = boxes[..., 2:] - boxes[..., :2]
                 min_tgt_size = torch.min(tgt_boxes_wh, dim=-1)[0]
-                keep = (min_tgt_size >= 1)
+                keep = (min_tgt_size >= 8)
 
                 tgt["boxes"] = boxes[keep]
                 tgt["labels"] = labels[keep]

yolo/models/yolov3/yolov3_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/yolov3/yolov3_fpn.py

@@ -62,6 +62,17 @@ class Yolov3FPN(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/yolov3/yolov3_head.py

@@ -111,7 +111,6 @@ class Yolov3DetHead(nn.Module):
         self.cls_head_dim = cfg.head_dim
         self.reg_head_dim = cfg.head_dim
 
-
     def forward(self, feats):
         """
             feats: List[(Tensor)] [[B, C, H, W], ...]

yolo/models/yolov3/yolov3_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/yolov5_backbone.py

@@ -15,6 +15,7 @@ pretrained_urls = {
     'x': None,
 }
 
+
 # --------------------- Yolov3's Backbone -----------------------
 ## Modified DarkNet
 class Yolov5Backbone(nn.Module):

yolo/models/yolov5/yolov5_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/yolov5_head.py

@@ -71,16 +71,6 @@ class DetHead(nn.Module):
         self.cls_feats = nn.Sequential(*cls_feats)
         self.reg_feats = nn.Sequential(*reg_feats)
 
-        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):
         """
             in_feats: (Tensor) [B, C, H, W]
@@ -111,6 +101,16 @@ class Yolov5DetHead(nn.Module):
         self.cls_head_dim = cfg.head_dim
         self.reg_head_dim = cfg.head_dim
 
+        # 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, feats):
         """

yolo/models/yolov5/yolov5_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/yolov5_pafpn.py

@@ -77,6 +77,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/yolov5/yolov5_pred.py

@@ -26,7 +26,6 @@ class DetPredLayer(nn.Module):
         self.obj_pred = nn.Conv2d(self.cls_dim, 1 * self.num_anchors, kernel_size=1)
         self.cls_pred = nn.Conv2d(self.cls_dim, num_classes * self.num_anchors, kernel_size=1)
         self.reg_pred = nn.Conv2d(self.reg_dim, 4 * self.num_anchors, kernel_size=1)                
-
         self.init_bias()
         
     def init_bias(self):