From 2540fd4c1c2d9186126a71b3eb681d3a0a11861e Mon Sep 17 00:00:00 2001 From: Jing Qiu <61612323+Laughing-q@users.noreply.github.com> Date: Fri, 4 Jul 2025 18:07:46 +0800 Subject: [PATCH] FIx downstream impact by the replacement of `torch.load` in `ultralytics` package (#13623) * Update torch_load * Auto-format by https://ultralytics.com/actions --------- Co-authored-by: UltralyticsAssistant --- export.py | 3 ++- hubconf.py | 4 ++-- models/experimental.py | 3 ++- pyproject.toml | 2 +- requirements.txt | 2 +- segment/train.py | 6 ++++-- train.py | 6 ++++-- utils/aws/resume.py | 3 ++- utils/general.py | 3 ++- 9 files changed, 20 insertions(+), 12 deletions(-) diff --git a/export.py b/export.py index 5e713e1..8ff652e 100644 --- a/export.py +++ b/export.py @@ -1128,10 +1128,11 @@ def pipeline_coreml(model, im, file, names, y, mlmodel, prefix=colorstr("CoreML Examples: ```python + from ultralytics.utils.patches import torch_load from pathlib import Path import torch - model = torch.load('yolov5s.pt') # Load YOLOv5 model + model = torch_load('yolov5s.pt') # Load YOLOv5 model im = torch.zeros((1, 3, 640, 640)) # Example input tensor names = {0: "person", 1: "bicycle", 2: "car", ...} # Define class names diff --git a/hubconf.py b/hubconf.py index 9d10630..5d00278 100644 --- a/hubconf.py +++ b/hubconf.py @@ -10,7 +10,7 @@ Usage: model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local') # local repo """ -import torch +from ultralytics.utils.patches import torch_load def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): @@ -87,7 +87,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo cfg = list((Path(__file__).parent / "models").rglob(f"{path.stem}.yaml"))[0] # model.yaml path model = DetectionModel(cfg, channels, classes) # create model if pretrained: - ckpt = torch.load(attempt_download(path), map_location=device) # load + ckpt = torch_load(attempt_download(path), map_location=device) # load csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=["anchors"]) # intersect model.load_state_dict(csd, strict=False) # load diff --git a/models/experimental.py b/models/experimental.py index 63d9c46..ab322e2 100644 --- a/models/experimental.py +++ b/models/experimental.py @@ -6,6 +6,7 @@ import math import numpy as np import torch import torch.nn as nn +from ultralytics.utils.patches import torch_load from utils.downloads import attempt_download @@ -95,7 +96,7 @@ def attempt_load(weights, device=None, inplace=True, fuse=True): model = Ensemble() for w in weights if isinstance(weights, list) else [weights]: - ckpt = torch.load(attempt_download(w), map_location="cpu") # load + ckpt = torch_load(attempt_download(w), map_location="cpu") # load ckpt = (ckpt.get("ema") or ckpt["model"]).to(device).float() # FP32 model # Model compatibility updates diff --git a/pyproject.toml b/pyproject.toml index 4e7a85e..fa91d79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,7 @@ dependencies = [ "thop>=0.1.1", # FLOPs computation "pandas>=1.1.4", "seaborn>=0.11.0", # plotting - "ultralytics>=8.1.47" + "ultralytics>=8.2.64" ] # Optional dependencies ------------------------------------------------------------------------------------------------ diff --git a/requirements.txt b/requirements.txt index dcd23bf..af57089 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,7 +15,7 @@ thop>=0.1.1 # FLOPs computation torch>=1.8.0 # see https://pytorch.org/get-started/locally (recommended) torchvision>=0.9.0 tqdm>=4.66.3 -ultralytics>=8.2.34 # https://ultralytics.com +ultralytics>=8.2.64 # https://ultralytics.com # protobuf<=3.20.1 # https://github.com/ultralytics/yolov5/issues/8012 # Logging --------------------------------------------------------------------- diff --git a/segment/train.py b/segment/train.py index 815c97c..bbea64f 100644 --- a/segment/train.py +++ b/segment/train.py @@ -40,6 +40,8 @@ if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATH ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative +from ultralytics.utils.patches import torch_load + import segment.val as validate # for end-of-epoch mAP from models.experimental import attempt_load from models.yolo import SegmentationModel @@ -174,7 +176,7 @@ def train(hyp, opt, device, callbacks): if pretrained: with torch_distributed_zero_first(LOCAL_RANK): weights = attempt_download(weights) # download if not found locally - ckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak + ckpt = torch_load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak model = SegmentationModel(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keys csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 @@ -605,7 +607,7 @@ def main(opt, callbacks=Callbacks()): with open(opt_yaml, errors="ignore") as f: d = yaml.safe_load(f) else: - d = torch.load(last, map_location="cpu")["opt"] + d = torch_load(last, map_location="cpu")["opt"] opt = argparse.Namespace(**d) # replace opt.cfg, opt.weights, opt.resume = "", str(last), True # reinstate if is_url(opt_data): diff --git a/train.py b/train.py index 1401ccb..5bb5e1f 100644 --- a/train.py +++ b/train.py @@ -44,6 +44,8 @@ if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATH ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative +from ultralytics.utils.patches import torch_load + import val as validate # for end-of-epoch mAP from models.experimental import attempt_load from models.yolo import Model @@ -212,7 +214,7 @@ def train(hyp, opt, device, callbacks): if pretrained: with torch_distributed_zero_first(LOCAL_RANK): weights = attempt_download(weights) # download if not found locally - ckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak + ckpt = torch_load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keys csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 @@ -646,7 +648,7 @@ def main(opt, callbacks=Callbacks()): with open(opt_yaml, errors="ignore") as f: d = yaml.safe_load(f) else: - d = torch.load(last, map_location="cpu")["opt"] + d = torch_load(last, map_location="cpu")["opt"] opt = argparse.Namespace(**d) # replace opt.cfg, opt.weights, opt.resume = "", str(last), True # reinstate if is_url(opt_data): diff --git a/utils/aws/resume.py b/utils/aws/resume.py index 5b80fd4..18fc030 100644 --- a/utils/aws/resume.py +++ b/utils/aws/resume.py @@ -9,6 +9,7 @@ from pathlib import Path import torch import yaml +from ultralytics.utils.patches import torch_load FILE = Path(__file__).resolve() ROOT = FILE.parents[2] # YOLOv5 root directory @@ -18,7 +19,7 @@ if str(ROOT) not in sys.path: port = 0 # --master_port path = Path("").resolve() for last in path.rglob("*/**/last.pt"): - ckpt = torch.load(last) + ckpt = torch_load(last) if ckpt["optimizer"] is None: continue diff --git a/utils/general.py b/utils/general.py index 89bbc61..efc0112 100644 --- a/utils/general.py +++ b/utils/general.py @@ -44,6 +44,7 @@ except (ImportError, AssertionError): import ultralytics from ultralytics.utils.checks import check_requirements +from ultralytics.utils.patches import torch_load from utils import TryExcept, emojis from utils.downloads import curl_download, gsutil_getsize @@ -1127,7 +1128,7 @@ def strip_optimizer(f="best.pt", s=""): Example: from utils.general import *; strip_optimizer() """ - x = torch.load(f, map_location=torch.device("cpu")) + x = torch_load(f, map_location=torch.device("cpu")) if x.get("ema"): x["model"] = x["ema"] # replace model with ema for k in "optimizer", "best_fitness", "ema", "updates": # keys