Check PyTorch model status for all YOLO methods (#945)

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>

ultralytics/yolo/utils/checks.py

@@ -234,17 +234,17 @@ def check_yolov5u_filename(file: str):
     return file
 
 
-def check_file(file, suffix=''):
+def check_file(file, suffix='', download=True):
     # Search/download file (if necessary) and return path
     check_suffix(file, suffix)  # optional
     file = str(file)  # convert to string
     file = check_yolov5u_filename(file)  # yolov5n -> yolov5nu
-    if not file or ('://' not in file and Path(file).is_file()):  # exists ('://' check required in Windows Python<3.10)
+    if not file or ('://' not in file and Path(file).exists()):  # exists ('://' check required in Windows Python<3.10)
         return file
-    elif file.lower().startswith(('https://', 'http://', 'rtsp://', 'rtmp://')):  # download
+    elif download and file.lower().startswith(('https://', 'http://', 'rtsp://', 'rtmp://')):  # download
         url = file  # warning: Pathlib turns :// -> :/
         file = Path(urllib.parse.unquote(file).split('?')[0]).name  # '%2F' to '/', split https://url.com/file.txt?auth
-        if Path(file).is_file():
+        if Path(file).exists():
             LOGGER.info(f'Found {url} locally at {file}')  # file already exists
         else:
             downloads.safe_download(url=url, file=file, unzip=False)

ultralytics/yolo/utils/dist.py

@@ -44,11 +44,17 @@ def generate_ddp_file(trainer):
 
 def generate_ddp_command(world_size, trainer):
     import __main__  # noqa local import to avoid https://github.com/Lightning-AI/lightning/issues/15218
-    file = generate_ddp_file(trainer) if sys.argv[0].endswith('yolo') else os.path.abspath(sys.argv[0])
+
+    # Get file and args (do not use sys.argv due to security vulnerability)
+    exclude_args = ['save_dir']
+    args = [f"{k}={v}" for k, v in vars(trainer.args).items() if k not in exclude_args]
+    file = generate_ddp_file(trainer)  # if argv[0].endswith('yolo') else os.path.abspath(argv[0])
+
+    # Build command
     torch_distributed_cmd = "torch.distributed.run" if TORCH_1_9 else "torch.distributed.launch"
     cmd = [
         sys.executable, "-m", torch_distributed_cmd, "--nproc_per_node", f"{world_size}", "--master_port",
-        f"{find_free_network_port()}", file] + sys.argv[1:]
+        f"{find_free_network_port()}", file] + args
     return cmd, file
 
 

ultralytics/yolo/utils/torch_utils.py

@@ -242,6 +242,11 @@ def copy_attr(a, b, include=(), exclude=()):
             setattr(a, k, v)
 
 
+def get_latest_opset():
+    # Return max supported ONNX opset by this version of torch
+    return max(int(k[14:]) for k in vars(torch.onnx) if 'symbolic_opset' in k)  # opset
+
+
 def intersect_dicts(da, db, exclude=()):
     # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
     return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}