Accept model=yolov5*u.yaml arguments (#4230)

ultralytics/utils/callbacks/dvc.py

@@ -1,4 +1,5 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
+
 import os
 
 import pkg_resources as pkg

ultralytics/utils/callbacks/wb.py

@@ -1,4 +1,5 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
+
 from ultralytics.utils import SETTINGS, TESTS_RUNNING
 from ultralytics.utils.torch_utils import model_info_for_loggers
 

ultralytics/utils/checks.py

@@ -341,15 +341,19 @@ def check_suffix(file='yolov8n.pt', suffix='.pt', msg=''):
 
 def check_yolov5u_filename(file: str, verbose: bool = True):
     """Replace legacy YOLOv5 filenames with updated YOLOv5u filenames."""
-    if ('yolov3' in file or 'yolov5' in file) and 'u' not in file:
-        original_file = file
-        file = re.sub(r'(.*yolov5([nsmlx]))\.pt', '\\1u.pt', file)  # i.e. yolov5n.pt -> yolov5nu.pt
-        file = re.sub(r'(.*yolov5([nsmlx])6)\.pt', '\\1u.pt', file)  # i.e. yolov5n6.pt -> yolov5n6u.pt
-        file = re.sub(r'(.*yolov3(|-tiny|-spp))\.pt', '\\1u.pt', file)  # i.e. yolov3-spp.pt -> yolov3-sppu.pt
-        if file != original_file and verbose:
-            LOGGER.info(f"PRO TIP 💡 Replace 'model={original_file}' with new 'model={file}'.\nYOLOv5 'u' models are "
-                        f'trained with https://github.com/ultralytics/ultralytics and feature improved performance vs '
-                        f'standard YOLOv5 models trained with https://github.com/ultralytics/yolov5.\n')
+    if 'yolov3' in file or 'yolov5' in file:
+        if 'u.yaml' in file:
+            file = file.replace('u.yaml', '.yaml')  # i.e. yolov5nu.yaml -> yolov5n.yaml
+        elif '.pt' in file and 'u' not in file:
+            original_file = file
+            file = re.sub(r'(.*yolov5([nsmlx]))\.pt', '\\1u.pt', file)  # i.e. yolov5n.pt -> yolov5nu.pt
+            file = re.sub(r'(.*yolov5([nsmlx])6)\.pt', '\\1u.pt', file)  # i.e. yolov5n6.pt -> yolov5n6u.pt
+            file = re.sub(r'(.*yolov3(|-tiny|-spp))\.pt', '\\1u.pt', file)  # i.e. yolov3-spp.pt -> yolov3-sppu.pt
+            if file != original_file and verbose:
+                LOGGER.info(
+                    f"PRO TIP 💡 Replace 'model={original_file}' with new 'model={file}'.\nYOLOv5 'u' models are "
+                    f'trained with https://github.com/ultralytics/ultralytics and feature improved performance vs '
+                    f'standard YOLOv5 models trained with https://github.com/ultralytics/yolov5.\n')
     return file
 
 

ultralytics/utils/torch_utils.py

@@ -1,4 +1,5 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
+
 import math
 import os
 import platform

ultralytics/utils/tuner.py

@@ -1,4 +1,5 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
+
 from ultralytics.cfg import TASK2DATA, TASK2METRIC
 from ultralytics.utils import DEFAULT_CFG_DICT, LOGGER, NUM_THREADS