fix(ml-backend): 推理分批,避免 Retrieve Predictions 撑爆显存;新增批量推送脚本

Retrieve Predictions 会把选中的任务一次发给后端,而后端把整个列表交给
model.predict(source=paths)——ultralytics 不理会 batch 参数,会当成一个 batch,
于是显存爆掉。这与之前 visualize_detections 踩的是同一个坑。后端现在按
YOLO_BATCH(默认 16)自己分批。

另加 scripts/push_predictions.py:重训完在本地分批推理,再用批量导入接口
POST /api/projects/<id>/import/predictions 提交,不经过后端服务,显存与请求体都可控。
支持 --replace 先删同版本旧预测(预测是追加的,重复跑会堆两份)、--limit 与 --dry-run。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-22 12:11:06 +08:00
co-authored by Claude Opus 5
parent 8fa7a9d142
commit dce2eef0a7
4 changed files with 374 additions and 3 deletions
+12 -3
View File
@@ -6,6 +6,8 @@
YOLO_CONF 置信度阈值,默认 0.25
YOLO_IMGSZ 推理尺寸,默认 640
YOLO_DEVICE cuda / cpu / mps,默认 cuda(不可用时 ultralytics 自己退回)
YOLO_BATCH 一次推理几张,默认 16。Retrieve Predictions 会一次发来成百上千个
任务,不分批会直接把显存撑爆
RETRAIN_AFTER 累计多少条新标注后自动触发重训,默认 0=只靠手动触发
RETRAIN_COMMAND 重训命令;不设则只记数不执行
RETRAIN_LOG 重训命令的输出写到哪,默认 /tmp/retrain.log
@@ -82,9 +84,16 @@ class FilmDefectDetector(LabelStudioMLBase):
if not paths:
return ModelResponse(model_version=model_version, predictions=[])
results = model.predict(
source=paths, imgsz=imgsz, conf=conf, device=device, verbose=False
)
# 必须自己分批:ultralytics 不理会 batch 参数,会把整个列表当成一个 batch。
# Label Studio 的 Retrieve Predictions 一次能发来上千个任务,不分批必然 OOM
batch = int(os.environ.get("YOLO_BATCH", 16))
results = []
for start in range(0, len(paths), batch):
results.extend(model.predict(
source=paths[start:start + batch], imgsz=imgsz, conf=conf,
device=device, verbose=False,
))
predictions = []
for result, (width, height) in zip(results, sizes, strict=True):
boxes = [