fix(detect): 可视化推理自己按批切分,避免 CUDA OOM

model.predict 不理会传入的 batch 参数,把整个图片列表当成一个 batch,8969 张图在
16 GiB 显卡上直接 OOM(单次要分配 10.27 GiB)。改成自己按 --batch 切片调用。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-17 21:47:44 +08:00
co-authored by Claude Opus 5
parent 788bc36ab8
commit 43c8c4a179
+11 -4
View File
@@ -159,12 +159,19 @@ def main() -> int:
rows = []
processed = 0
for split, images in splits.items():
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 images], stream=True, imgsz=args.imgsz, conf=args.conf,
device=device_arg, batch=args.batch, verbose=False,
source=[str(p) for p in chunk], imgsz=args.imgsz, conf=args.conf,
device=device_arg, verbose=False,
)
for image_path, result in zip(images, results, strict=True):
yield from zip(chunk, results, strict=True)
for split, images in splits.items():
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(