Update models/hub/*.yaml files for v6.0n release (#5540)

* Update model yamls for v6.0 * Add python models/yolo.py --test * Ghost fix
ultralytics · Nov 6, 2021 · fa2344c · fa2344c
1 parent 76d90d8
commit fa2344c
Show file tree

Hide file tree

Showing 9 changed files with 72 additions and 63 deletions.
diff --git a/models/hub/yolov5-bifpn.yaml b/models/hub/yolov5-bifpn.yaml
@@ -9,22 +9,22 @@ anchors:
   - [30,61, 62,45, 59,119]  # P4/16
   - [116,90, 156,198, 373,326]  # P5/32
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
    [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, C3, [256]],
+   [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
    [-1, 9, C3, [512]],
    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
-   [-1, 1, SPP, [1024, [5, 9, 13]]],
-   [-1, 3, C3, [1024, False]],  # 9
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 9
   ]
 
-# YOLOv5 BiFPN head
+# YOLOv5 v6.0 BiFPN head
 head:
   [[-1, 1, Conv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
@@ -37,7 +37,7 @@ head:
    [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
 
    [-1, 1, Conv, [256, 3, 2]],
-   [[-1, 14, 6], 1, Concat, [1]],  # cat P4
+   [[-1, 14, 6], 1, Concat, [1]],  # cat P4 <--- BiFPN change
    [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
 
    [-1, 1, Conv, [512, 3, 2]],

diff --git a/models/hub/yolov5-fpn.yaml b/models/hub/yolov5-fpn.yaml
@@ -9,34 +9,34 @@ anchors:
   - [30,61, 62,45, 59,119]  # P4/16
   - [116,90, 156,198, 373,326]  # P5/32
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
-   [-1, 3, Bottleneck, [128]],
+   [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, BottleneckCSP, [256]],
+   [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
-   [-1, 9, BottleneckCSP, [512]],
+   [-1, 9, C3, [512]],
    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
-   [-1, 1, SPP, [1024, [5, 9, 13]]],
-   [-1, 6, BottleneckCSP, [1024]],  # 9
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 9
   ]
 
-# YOLOv5 FPN head
+# YOLOv5 v6.0 FPN head
 head:
-  [[-1, 3, BottleneckCSP, [1024, False]],  # 10 (P5/32-large)
+  [[-1, 3, C3, [1024, False]],  # 10 (P5/32-large)
 
    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
    [[-1, 6], 1, Concat, [1]],  # cat backbone P4
    [-1, 1, Conv, [512, 1, 1]],
-   [-1, 3, BottleneckCSP, [512, False]],  # 14 (P4/16-medium)
+   [-1, 3, C3, [512, False]],  # 14 (P4/16-medium)
 
    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
    [[-1, 4], 1, Concat, [1]],  # cat backbone P3
    [-1, 1, Conv, [256, 1, 1]],
-   [-1, 3, BottleneckCSP, [256, False]],  # 18 (P3/8-small)
+   [-1, 3, C3, [256, False]],  # 18 (P3/8-small)
 
    [[18, 14, 10], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
   ]
diff --git a/models/hub/yolov5-p2.yaml b/models/hub/yolov5-p2.yaml
@@ -4,24 +4,24 @@
 nc: 80  # number of classes
 depth_multiple: 1.0  # model depth multiple
 width_multiple: 1.0  # layer channel multiple
-anchors: 3
+anchors: 3  # auto-anchor evolves 3 anchors per P output layer
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
    [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, C3, [256]],
+   [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
    [-1, 9, C3, [512]],
    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
-   [-1, 1, SPP, [1024, [5, 9, 13]]],
-   [-1, 3, C3, [1024, False]],  # 9
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 9
   ]
 
-# YOLOv5 head
+# YOLOv5 v6.0 head
 head:
   [[-1, 1, Conv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, 'nearest']],

diff --git a/models/hub/yolov5-p6.yaml b/models/hub/yolov5-p6.yaml
@@ -4,26 +4,26 @@
 nc: 80  # number of classes
 depth_multiple: 1.0  # model depth multiple
 width_multiple: 1.0  # layer channel multiple
-anchors: 3
+anchors: 3  # auto-anchor 3 anchors per P output layer
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
    [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, C3, [256]],
+   [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
    [-1, 9, C3, [512]],
    [-1, 1, Conv, [768, 3, 2]],  # 7-P5/32
    [-1, 3, C3, [768]],
    [-1, 1, Conv, [1024, 3, 2]],  # 9-P6/64
-   [-1, 1, SPP, [1024, [3, 5, 7]]],
-   [-1, 3, C3, [1024, False]],  # 11
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 11
   ]
 
-# YOLOv5 head
+# YOLOv5 v6.0 head
 head:
   [[-1, 1, Conv, [768, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
@@ -50,7 +50,7 @@ head:
 
    [-1, 1, Conv, [768, 3, 2]],
    [[-1, 12], 1, Concat, [1]],  # cat head P6
-   [-1, 3, C3, [1024, False]],  # 32 (P5/64-xlarge)
+   [-1, 3, C3, [1024, False]],  # 32 (P6/64-xlarge)
 
    [[23, 26, 29, 32], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5, P6)
   ]
diff --git a/models/hub/yolov5-p7.yaml b/models/hub/yolov5-p7.yaml
@@ -4,25 +4,25 @@
 nc: 80  # number of classes
 depth_multiple: 1.0  # model depth multiple
 width_multiple: 1.0  # layer channel multiple
-anchors: 3
+anchors: 3  # auto-anchor 3 anchors per P output layer
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
    [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, C3, [256]],
+   [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
    [-1, 9, C3, [512]],
    [-1, 1, Conv, [768, 3, 2]],  # 7-P5/32
    [-1, 3, C3, [768]],
    [-1, 1, Conv, [1024, 3, 2]],  # 9-P6/64
    [-1, 3, C3, [1024]],
    [-1, 1, Conv, [1280, 3, 2]],  # 11-P7/128
-   [-1, 1, SPP, [1280, [3, 5]]],
-   [-1, 3, C3, [1280, False]],  # 13
+   [-1, 3, C3, [1280]],
+   [-1, 1, SPPF, [1280, 5]],  # 13
   ]
 
 # YOLOv5 head

diff --git a/models/hub/yolov5-panet.yaml b/models/hub/yolov5-panet.yaml
@@ -9,40 +9,40 @@ anchors:
   - [30,61, 62,45, 59,119]  # P4/16
   - [116,90, 156,198, 373,326]  # P5/32
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
-   [-1, 3, BottleneckCSP, [128]],
+   [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, BottleneckCSP, [256]],
+   [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
-   [-1, 9, BottleneckCSP, [512]],
+   [-1, 9, C3, [512]],
    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
-   [-1, 1, SPP, [1024, [5, 9, 13]]],
-   [-1, 3, BottleneckCSP, [1024, False]],  # 9
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 9
   ]
 
-# YOLOv5 PANet head
+# YOLOv5 v6.0 PANet 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, BottleneckCSP, [512, False]],  # 13
+   [-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, BottleneckCSP, [256, False]],  # 17 (P3/8-small)
+   [-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, BottleneckCSP, [512, False]],  # 20 (P4/16-medium)
+   [-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, BottleneckCSP, [1024, False]],  # 23 (P5/32-large)
+   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
 
    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
   ]
diff --git a/models/hub/yolov5s-ghost.yaml b/models/hub/yolov5s-ghost.yaml
@@ -9,22 +9,22 @@ anchors:
   - [30,61, 62,45, 59,119]  # P4/16
   - [116,90, 156,198, 373,326]  # P5/32
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, GhostConv, [128, 3, 2]],  # 1-P2/4
    [-1, 3, C3Ghost, [128]],
    [-1, 1, GhostConv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, C3Ghost, [256]],
+   [-1, 6, C3Ghost, [256]],
    [-1, 1, GhostConv, [512, 3, 2]],  # 5-P4/16
    [-1, 9, C3Ghost, [512]],
    [-1, 1, GhostConv, [1024, 3, 2]],  # 7-P5/32
-   [-1, 1, SPP, [1024, [5, 9, 13]]],
-   [-1, 3, C3Ghost, [1024, False]],  # 9
+   [-1, 3, C3Ghost, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 9
   ]
 
-# YOLOv5 head
+# YOLOv5 v6.0 head
 head:
   [[-1, 1, GhostConv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, 'nearest']],

diff --git a/models/hub/yolov5s-transformer.yaml b/models/hub/yolov5s-transformer.yaml
@@ -9,22 +9,22 @@ anchors:
   - [30,61, 62,45, 59,119]  # P4/16
   - [116,90, 156,198, 373,326]  # P5/32
 
-# YOLOv5 backbone
+# YOLOv5 v6.0 backbone
 backbone:
   # [from, number, module, args]
-  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
    [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
-   [-1, 9, C3, [256]],
+   [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
    [-1, 9, C3, [512]],
    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
-   [-1, 1, SPP, [1024, [5, 9, 13]]],
-   [-1, 3, C3TR, [1024, False]],  # 9  <-------- C3TR() Transformer module
+   [-1, 3, C3TR, [1024]],  # 9 <--- C3TR() Transformer module
+   [-1, 1, SPPF, [1024, 5]],  # 9
   ]
 
-# YOLOv5 head
+# YOLOv5 v6.0 head
 head:
   [[-1, 1, Conv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, 'nearest']],

diff --git a/models/yolo.py b/models/yolo.py
@@ -306,6 +306,7 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
     parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
     parser.add_argument('--profile', action='store_true', help='profile model speed')
+    parser.add_argument('--test', action='store_true', help='test all yolo*.yaml')
     opt = parser.parse_args()
     opt.cfg = check_yaml(opt.cfg)  # check YAML
     print_args(FILE.stem, opt)
@@ -320,6 +321,14 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
         img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
         y = model(img, profile=True)
 
+    # Test all models
+    if opt.test:
+        for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'):
+            try:
+                _ = Model(cfg)
+            except Exception as e:
+                print(f'Error in {cfg}: {e}')
+
     # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
     # from torch.utils.tensorboard import SummaryWriter
     # tb_writer = SummaryWriter('.')