From 43c8c4a1798b22787e186ecf426f3daf7fba3859 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Thu, 17 Sep 2026 21:47:44 +0800 Subject: [PATCH] =?UTF-8?q?fix(detect):=20=E5=8F=AF=E8=A7=86=E5=8C=96?= =?UTF-8?q?=E6=8E=A8=E7=90=86=E8=87=AA=E5=B7=B1=E6=8C=89=E6=89=B9=E5=88=87?= =?UTF-8?q?=E5=88=86=EF=BC=8C=E9=81=BF=E5=85=8D=20CUDA=20OOM?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit model.predict 不理会传入的 batch 参数，把整个图片列表当成一个 batch，8969 张图在 16 GiB 显卡上直接 OOM（单次要分配 10.27 GiB）。改成自己按 --batch 切片调用。 Co-Authored-By: Claude Opus 5 --- scripts/visualize_detections.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/scripts/visualize_detections.py b/scripts/visualize_detections.py index 02e19c5..a4d8a3a 100644 --- a/scripts/visualize_detections.py +++ b/scripts/visualize_detections.py @@ -159,12 +159,19 @@ def main() -> int: rows = [] processed = 0 + def predict_in_batches(images): + """自己按批切分：ultralytics 的 predict 不理会 batch 参数， + 一次性把整个列表当成一个 batch，上万张图必然 OOM。""" + for start in range(0, len(images), args.batch): + chunk = images[start:start + args.batch] + results = model.predict( + source=[str(p) for p in chunk], imgsz=args.imgsz, conf=args.conf, + device=device_arg, verbose=False, + ) + yield from zip(chunk, results, strict=True) + for split, images in splits.items(): - results = model.predict( - source=[str(p) for p in images], stream=True, imgsz=args.imgsz, conf=args.conf, - device=device_arg, batch=args.batch, verbose=False, - ) - for image_path, result in zip(images, results, strict=True): + for image_path, result in predict_in_batches(images): predictions = [ Box(int(cls), *(float(v) for v in xyxy), float(score)) for xyxy, cls, score in zip(