浏览代码

train YOLOv5-S

yjh0410 2 年之前
父节点
当前提交
22861e1213

+ 1 - 1
README.md

@@ -108,7 +108,7 @@ python train.py --cuda -d coco --root path/to/COCO -v yolov1 -bs 16 --max_epoch
 
 | Model         |   Backbone         | Scale | Epoch |  FPS  | AP<sup>val<br>0.5:0.95 | AP<sup>val<br>0.5 | FLOPs<br><sup>(G) | Params<br><sup>(M) | Weight |
 |---------------|--------------------|-------|-------|-------|------------------------|-------------------|-------------------|--------------------|--------|
-| YOLOv3-Tiny   | DarkNet-Tiny       |  640  |  250  |       |                        |                   |                   |                    |  |
+| YOLOv3-Tiny   | DarkNet-Tiny       |  640  |  250  |       |                        |                   |   7.0             |   2.3              |  |
 | YOLOv3        | DarkNet-53         |  640  |  250  |       |        42.9            |       63.5        |   167.4           |   54.9             | [ckpt](https://github.com/yjh0410/PyTorch_YOLO_Tutorial/releases/download/yolo_tutorial_ckpt/yolov3_coco.pth) |
 
 * YOLOv4:

+ 1 - 1
models/detectors/yolov3/yolov3.py

@@ -203,7 +203,7 @@ class YOLOv3(nn.Module):
         return bboxes, scores, labels
 
 
-
+    # ---------------------- Main Process for Training ----------------------
     def forward(self, x):
         if not self.trainable:
             return self.inference(x)

+ 2 - 2
models/detectors/yolov3/yolov3_backbone.py

@@ -8,7 +8,7 @@ except:
     
 
 model_urls = {
-    "darknet_tiny": None,
+    "darknet_tiny": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/darknet_tiny.pth",
     "darknet53": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/darknet53_silu.pth",
 }
 
@@ -148,7 +148,7 @@ def build_backbone(model_name='darknet53', pretrained=False):
 if __name__ == '__main__':
     import time
     from thop import profile
-    model, feats = build_backbone(pretrained=False)
+    model, feats = build_backbone(model_name='darknet_tiny', pretrained=True)
     x = torch.randn(1, 3, 224, 224)
     t0 = time.time()
     outputs = model(x)

+ 2 - 2
models/detectors/yolov4/yolov4.py

@@ -5,7 +5,7 @@ from utils.misc import multiclass_nms
 
 from .yolov4_backbone import build_backbone
 from .yolov4_neck import build_neck
-from .yolov4_fpn import build_fpn
+from .yolov4_pafpn import build_fpn
 from .yolov4_head import build_head
 
 
@@ -203,7 +203,7 @@ class YOLOv4(nn.Module):
         return bboxes, scores, labels
 
 
-
+    # ---------------------- Main Process for Training ----------------------
     def forward(self, x):
         if not self.trainable:
             return self.inference(x)

+ 0 - 0
models/detectors/yolov4/yolov4_fpn.py → models/detectors/yolov4/yolov4_pafpn.py


+ 1 - 1
models/detectors/yolov5/yolov5_backbone.py

@@ -11,7 +11,7 @@ except:
 model_urls = {
     "cspdarknet_nano": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/cspdarknet_nano.pth",
     "cspdarknet_small": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/cspdarknet_small.pth",
-    "cspdarknet_medium": None,
+    "cspdarknet_medium": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/cspdarknet_medium.pth",
     "cspdarknet_large": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/cspdarknet_large.pth",
     "cspdarknet_huge": None,
 }

+ 1 - 1
models/detectors/yolov7/yolov7.py

@@ -5,7 +5,7 @@ from utils.misc import multiclass_nms
 
 from .yolov7_backbone import build_backbone
 from .yolov7_neck import build_neck
-from .yolov7_fpn import build_fpn
+from .yolov7_pafpn import build_fpn
 from .yolov7_head import build_head
 
 

+ 0 - 0
models/detectors/yolov7/yolov7_fpn.py → models/detectors/yolov7/yolov7_pafpn.py


+ 1 - 1
train.sh

@@ -3,7 +3,7 @@ python train.py \
         --cuda \
         -d coco \
         --root /mnt/share/ssd2/dataset/ \
-        -m yolov7_l \
+        -m yolov5_s \
         -bs 16 \
         -size 640 \
         --wp_epoch 1 \