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:
co-authored by
Claude Opus 5
parent
788bc36ab8
commit
43c8c4a179
@@ -159,12 +159,19 @@ def main() -> int:
|
|||||||
rows = []
|
rows = []
|
||||||
processed = 0
|
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(
|
results = model.predict(
|
||||||
source=[str(p) for p in images], stream=True, imgsz=args.imgsz, conf=args.conf,
|
source=[str(p) for p in chunk], imgsz=args.imgsz, conf=args.conf,
|
||||||
device=device_arg, batch=args.batch, verbose=False,
|
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 = [
|
predictions = [
|
||||||
Box(int(cls), *(float(v) for v in xyxy), float(score))
|
Box(int(cls), *(float(v) for v in xyxy), float(score))
|
||||||
for xyxy, cls, score in zip(
|
for xyxy, cls, score in zip(
|
||||||
|
|||||||
Reference in New Issue
Block a user