xzhou/ModelTrainEval
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ml-backend): Start Training 按钮真正触发重训，启动/重训脚本进 scripts/

Browse Source 
原先点 Label Studio 的 Start Training 只是把计数加一：START_TRAINING 和普通标注事件
走同一条阈值判断，没配 RETRAIN_COMMAND 时更是静默什么都不做。实际使用中用户点了按钮、
以为在训练，其实权重一直没变。

- Start Training 无条件触发重训；攒够 RETRAIN_AFTER 条新标注也自动触发。
- 只数 ANNOTATION_CREATED：UPDATED / DELETED 是对已有标注的修改，不该把阈值顶到。
- 已有重训在跑时重复点击被忽略，不叠加启动。
- 没配 RETRAIN_COMMAND 时打警告说明不会真的训练，并在返回值里说清楚。
- 子进程输出重定向到 RETRAIN_LOG（默认 /tmp/retrain.log），训练进度有地方看。

scripts/start_ml_backend.sh、scripts/retrain_from_label_studio.sh 放进仓库，
可从工作区直接启动。重训末尾把 runs/detect/current.pt 软链接指向新权重，
启动脚本默认读它，换模型不必改脚本。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-21 11:00:01 +08:00
co-authored by Claude Opus 5
parent 1b8c8d7cac
commit e61fd5887c
5 changed files with 211 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml_backend/model.py
+43 -13
View File
@@ -6,8 +6,9 @@
YOLO_CONF 置信度阈值，默认 0.25
YOLO_IMGSZ 推理尺寸，默认 640
YOLO_DEVICE cuda / cpu / mps，默认 cuda（不可用时 ultralytics 自己退回）
RETRAIN_AFTER 累计多少条新标注后触发重训，默认 0=不触发
RETRAIN_COMMAND 触发时执行的命令；不设则只记数不执行
RETRAIN_AFTER 累计多少条新标注后自动触发重训，默认 0=只靠手动触发
RETRAIN_COMMAND 重训命令；不设则只记数不执行
RETRAIN_LOG 重训命令的输出写到哪，默认 /tmp/retrain.log
"""
from __future__ import annotations
@@ -31,6 +32,7 @@ class FilmDefectDetector(LabelStudioMLBase):
_model = None
_lock = threading.Lock()
_retraining = None # 正在跑的重训子进程
def setup(self) -> None:
"""记录模型版本。
@@ -100,25 +102,53 @@ class FilmDefectDetector(LabelStudioMLBase):
return ModelResponse(model_version=model_version, predictions=predictions)
def fit(self, event, data, **kwargs):
"""每产生一条标注就被调用一次，用来攒够数量后触发重训。
"""Label Studio 的标注事件与 Start Training 按钮都走这里。
这里**不能**直接训练：Label Studio 等着这个请求返回，长任务会把它卡住。
所以只累加计数，达到阈值后起一个子进程去跑，立刻返回。
所以只决定"要不要训"，真正的训练交给子进程，立刻返回。
"""
if event not in self.TRAIN_EVENTS:
return None
after = int(os.environ.get("RETRAIN_AFTER", 0))
command = os.environ.get("RETRAIN_COMMAND", "")
count = int(self.get("annotations_since_training") or 0) + 1
self.set("annotations_since_training", str(count))
logger.info("收到 %s，距上次训练已有 %d 条标注", event, count)
count = int(self.get("annotations_since_training") or 0)
if event == "ANNOTATION_CREATED":
# 只数新增：UPDATED / DELETED 是对已有标注的修改，不该让阈值提前到点
count += 1
self.set("annotations_since_training", str(count))
logger.info("收到 %s，距上次训练已新增 %d 条标注", event, count)
if not after or count < after or not command:
# Start Training 是人明确点的，无条件训；攒够阈值则自动训
manual = event == "START_TRAINING"
if not manual and not (after and count >= after):
return {"annotations_since_training": count}
return self._start_retraining(count, manual)
def _start_retraining(self, count: int, manual: bool) -> dict:
"""起一个子进程跑重训命令，立刻返回。"""
command = os.environ.get("RETRAIN_COMMAND", "")
if not command:
logger.warning(
"%s，但没有设置 RETRAIN_COMMAND，不会真的训练"
"（在 scripts/start_ml_backend.sh 里配置）",
"收到 Start Training" if manual else f"已累计 {count} 条标注",
)
return {"annotations_since_training": count, "retrain_started": False,
"detail": "RETRAIN_COMMAND 未设置"}
running = FilmDefectDetector._retraining
if running is not None and running.poll() is None:
# 重复点 Start Training 不该叠加启动多个训练
logger.warning("已有重训在跑（pid=%s），这次忽略", running.pid)
return {"annotations_since_training": count, "retrain_started": False,
"detail": f"已有重训在跑 pid={running.pid}"}
log_path = os.environ.get("RETRAIN_LOG", "/tmp/retrain.log")
self.set("annotations_since_training", "0")
logger.info("累计 %d 条标注，触发重训: %s", count, command)
# 不等它结束：训练要几十分钟，这个请求必须立刻返回
subprocess.Popen(shlex.split(command)) # noqa: S603
return {"annotations_since_training": 0, "retrain_started": True}
logger.info("触发重训: %s（输出见 %s）", command, log_path)
with open(log_path, "a", encoding="utf-8") as log: # noqa: PTH123
FilmDefectDetector._retraining = subprocess.Popen( # noqa: S603
shlex.split(command), stdout=log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL,
)
return {"annotations_since_training": 0, "retrain_started": True,
"pid": FilmDefectDetector._retraining.pid, "log": log_path}