浏览代码

release YOLOv7-Tiny on COCO

yjh0410 2 年之前
父节点
当前提交
b6c5812a47
共有 2 个文件被更改,包括 3 次插入3 次删除
  1. 1 1
      README.md
  2. 2 2
      train.py

+ 1 - 1
README.md

@@ -107,7 +107,7 @@ python train.py --cuda -d coco --root path/to/COCO -v yolov1 -bs 16 --max_epoch
 | YOLOv5        | CSPDarkNet-53      |  640  |  √   |  250  |       |                        |                   |  |
 | YOLOX         | CSPDarkNet-L       |  640  |  √   |  300  |       |        46.6            |       66.1        | [ckpt](https://github.com/yjh0410/PyTorch_YOLO_Tutorial/releases/download/yolo_tutorial_ckpt/yolox_coco.pth) |
 | YOLOv7-Nano   | ELANNet-Nano       |  640  |  √   |  300  |       |                        |                   |  |
-| YOLOv7-Tiny   | ELANNet-Tiny       |  640  |  √   |  300  |       |                        |                   |  |
+| YOLOv7-Tiny   | ELANNet-Tiny       |  640  |  √   |  300  |       |        37.7            |       56.6        | [ckpt](https://github.com/yjh0410/PyTorch_YOLO_Tutorial/releases/download/yolo_tutorial_ckpt/yolov7_tiny_coco.pth) |
 | YOLOv7-Large  | ELANNet-Large      |  640  |  √   |  300  |       |                        |                   |  |
 | YOLOv8-Nano   | CSP-ELANNet-Nano   |  640  |  √   |  300  |       |                        |                   |  |
 | YOLOv8-Small  | CSP-ELANNet-Small  |  640  |  √   |  300  |       |                        |                   |  |

+ 2 - 2
train.py

@@ -182,7 +182,8 @@ def train():
     optimizer, start_epoch = build_optimizer(model_cfg, model_without_ddp, model_cfg['lr0'], args.resume)
 
     # Scheduler
-    scheduler, lf = build_lr_scheduler(model_cfg, optimizer, args.max_epoch)
+    total_epochs = args.max_epoch + args.wp_epoch
+    scheduler, lf = build_lr_scheduler(model_cfg, optimizer, total_epochs)
     scheduler.last_epoch = start_epoch - 1  # do not move
     if args.resume:
         scheduler.step()
@@ -197,7 +198,6 @@ def train():
     # start training loop
     best_map = -1.0
     last_opt_step = -1
-    total_epochs = args.max_epoch + args.wp_epoch
     heavy_eval = False
     optimizer.zero_grad()