From e61fd5887c262ff1bd205c9ab8034f1952ff4d88 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Mon, 21 Sep 2026 11:00:01 +0800 Subject: [PATCH] =?UTF-8?q?feat(ml-backend):=20Start=20Training=20?= =?UTF-8?q?=E6=8C=89=E9=92=AE=E7=9C=9F=E6=AD=A3=E8=A7=A6=E5=8F=91=E9=87=8D?= =?UTF-8?q?=E8=AE=AD=EF=BC=8C=E5=90=AF=E5=8A=A8/=E9=87=8D=E8=AE=AD?= =?UTF-8?q?=E8=84=9A=E6=9C=AC=E8=BF=9B=20scripts/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 原先点 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 --- ml_backend/README.md | 46 ++++++++++++---------- ml_backend/model.py | 56 ++++++++++++++++++++------- scripts/retrain_from_label_studio.sh | 50 ++++++++++++++++++++++++ scripts/start_ml_backend.sh | 58 ++++++++++++++++++++++++++++ tests/test_ml_backend_convert.py | 34 ++++++++++++++++ 5 files changed, 211 insertions(+), 33 deletions(-) create mode 100755 scripts/retrain_from_label_studio.sh create mode 100755 scripts/start_ml_backend.sh diff --git a/ml_backend/README.md b/ml_backend/README.md index 6bd7b5e..800c1b3 100644 --- a/ml_backend/README.md +++ b/ml_backend/README.md @@ -13,23 +13,26 @@ jinja2 3.1。装进主环境会直接弄坏 Label Studio。它本来就是个独 ```bash uv venv .venv-ml --python 3.12 +# --no-deps 装本体：它把 label-studio-sdk pin 到了 GitHub 某个提交， +# 那个仓库大、容易传输中断（服务器上实测失败过），依赖全部走 PyPI 更稳 +VIRTUAL_ENV=.venv-ml uv pip install --no-deps \ + "label-studio-ml @ git+https://github.com/HumanSignal/label-studio-ml-backend.git" VIRTUAL_ENV=.venv-ml uv pip install \ - "label-studio-ml @ git+https://github.com/HumanSignal/label-studio-ml-backend.git" \ - ultralytics + "flask>2.3" "colorama~=0.4" "requests~=2.31" "semver~=3.0.2" \ + "label-studio-sdk>=2.1" ultralytics ``` ## 启动 ```bash -cd ml_backend -YOLO_WEIGHTS=$PWD/../runs/detect/film/weights/best.pt \ -YOLO_DEVICE=cuda \ -LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT=/home/deep/labeling \ -../.venv-ml/bin/python _wsgi.py -p 9090 +bash scripts/start_ml_backend.sh ``` -**用 `python _wsgi.py` 而不是 `label-studio-ml start .`**：后者用 PATH 上的 python 起子进程， -在独立环境里会 `ModuleNotFoundError`。 +脚本会停掉旧进程、带齐环境变量重启，并打印权重、日志位置与网页端要填的 URL。 +所有参数都能用环境变量覆盖，例如 `YOLO_WEIGHTS=... bash scripts/start_ml_backend.sh`。 + +它内部用 `python _wsgi.py` 而不是 `label-studio-ml start .`：后者用 PATH 上的 python +起子进程，在独立环境里会 `ModuleNotFoundError`。 `LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 必须和启动 Label Studio 时的值一致——任务里的图片是 `/data/local-files/?d=...`，后端要靠它把 URL 还原成磁盘路径。图片若来自上传或云存储， @@ -66,19 +69,22 @@ Label Studio 里：Settings → **Model** → Connect Model，URL 填 `http://lo ## 边标边训 -Label Studio 每产生一条标注就会 POST `/webhook`，后端在 `fit()` 里累加计数， -达到 `RETRAIN_AFTER` 就起一个**子进程**执行 `RETRAIN_COMMAND` 并立刻返回—— -训练要几十分钟，绝不能在这个请求里同步做，否则会把 Label Studio 卡住。 +**在 Label Studio 里点 Settings → Model → Start Training**，后端收到 `START_TRAINING` +事件，起一个子进程跑 `scripts/retrain_from_label_studio.sh`，立刻返回——训练要几十分钟， +绝不能在这个请求里同步做，否则会把 Label Studio 卡住。 -```bash -RETRAIN_AFTER=200 \ -RETRAIN_COMMAND="/home/deep/ModelTrainEval/scripts/retrain.sh" \ -../.venv-ml/bin/python _wsgi.py -p 9090 -``` +触发规则： -`retrain.sh` 需要自己写，大致是：导出项目 → `build_yolo_dataset.py` 装配 → `detect.py` 训练 → -把新权重放到 `YOLO_WEIGHTS` 指向的位置 → 重启本服务（或让它下次启动时加载新权重）。 -**目前权重只在进程启动时加载一次**，重训完要重启服务才会生效。 +- **Start Training 按钮无条件触发**，不看攒了多少条； +- 新增标注累计到 `RETRAIN_AFTER` 也会自动触发（默认 0 = 只手动）。只数 + `ANNOTATION_CREATED`，`UPDATED` / `DELETED` 是对已有标注的修改，不该把阈值顶到； +- 已有重训在跑时重复点击会被忽略，不会叠起多个训练。 + +重训脚本做四件事：导出标注 → `build_yolo_dataset.py` 装配 → `detect.py` 训练 → +把 `runs/detect/current.pt` 软链接指向新权重并重启本服务。启动脚本默认读这个软链接， +所以换模型不用改任何脚本。**权重只在进程启动时加载一次**，这也是重训末尾要重启服务的原因。 + +训练进度看 `$RETRAIN_LOG`（默认 `/tmp/retrain.log`）；Label Studio 网页端不显示训练进度。 ## 已验证的行为 diff --git a/ml_backend/model.py b/ml_backend/model.py index 61bde84..506ac5a 100644 --- a/ml_backend/model.py +++ b/ml_backend/model.py @@ -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} diff --git a/scripts/retrain_from_label_studio.sh b/scripts/retrain_from_label_studio.sh new file mode 100755 index 0000000..6f8eec8 --- /dev/null +++ b/scripts/retrain_from_label_studio.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# 从 Label Studio 的最新标注重新训练检测模型，并让 ML backend 用上新权重。 +# +# bash scripts/retrain_from_label_studio.sh +# +# 也由 ML backend 在点击 Start Training 时自动调用（输出见 $RETRAIN_LOG）。 +# 四步：导出标注 -> 装配数据集 -> 训练 -> 切换权重并重启后端。 +set -euo pipefail + +REPO=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) +cd "$REPO" + +PROJECT="${LS_PROJECT_ID:-39}" +IMAGES="${LS_IMAGES_DIR:-$REPO/runs/boxes_shap_label_studio/images}" +STAMP=$(date +%Y%m%d_%H%M) +EXPORT="${LS_EXPORT_DIR:-/tmp}/ls_export_${STAMP}.zip" +DATASET="$REPO/data/film_yolo_${STAMP}" +WORK_DIR="$REPO/runs/detect/film_${STAMP}" + +PYTHON="${REPO_PYTHON:-$REPO/.venv/bin/python}" +# 用跑着 Label Studio 的那个环境导出，确保读的是同一个库 +LS_BIN="${LS_BIN:-/home/deep/.v/label/bin/label-studio}" +[ -x "$LS_BIN" ] || LS_BIN="$REPO/.venv/bin/label-studio" +export LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT="${LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT:-/home/deep/labeling}" + +# 训练参数的额外覆盖，冒烟测试时用，如 RETRAIN_EXTRA_SET="train.epochs=1 train.fraction=0.02" +EXTRA_SET="${RETRAIN_EXTRA_SET:-}" + +echo "===== [$(date '+%F %T')] 开始重训（项目 $PROJECT）=====" + +echo "--- 1/4 导出标注 -> $EXPORT" +"$LS_BIN" export "$PROJECT" YOLO --export-path="$EXPORT" + +echo "--- 2/4 装配数据集 -> $DATASET" +"$PYTHON" scripts/build_yolo_dataset.py \ + --export "$EXPORT" --images "$IMAGES" --out "$DATASET" --overwrite + +echo "--- 3/4 训练 -> $WORK_DIR" +# shellcheck disable=SC2086 # EXTRA_SET 需要按空格拆成多个参数 +"$PYTHON" detect.py --config configs/detect/film.yaml \ + --work-dir "$WORK_DIR" \ + --set data.yaml="$DATASET/data.yaml" $EXTRA_SET + +echo "--- 4/4 切换权重并重启 ML backend" +# 指向最新权重的稳定软链接：启动脚本读它，于是换模型不用改任何脚本 +ln -sfn "$WORK_DIR/weights/best.pt" "$REPO/runs/detect/current.pt" +bash "$REPO/scripts/start_ml_backend.sh" + +echo "===== [$(date '+%F %T')] 重训完成：$WORK_DIR =====" +echo "指标见 $WORK_DIR/metrics.json，曲线见 $WORK_DIR/results.png" diff --git a/scripts/start_ml_backend.sh b/scripts/start_ml_backend.sh new file mode 100755 index 0000000..991968f --- /dev/null +++ b/scripts/start_ml_backend.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +# 启动 Label Studio 的 ML backend（预标注 + 触发重训）。 +# +# bash scripts/start_ml_backend.sh +# +# 会先停掉旧进程再起新的。所有参数都可以用环境变量覆盖，例如换权重： +# +# YOLO_WEIGHTS=/path/to/best.pt bash scripts/start_ml_backend.sh +# +# 权重只在进程启动时加载一次，重训完必须重跑本脚本才会生效。 +set -euo pipefail + +REPO=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) + +# 默认用 runs/detect/current.pt 这个软链接：重训脚本训完会把它指向最新权重， +# 这样换模型不用改脚本。链接不存在时退回第一版训练的结果。 +DEFAULT_WEIGHTS="$REPO/runs/detect/current.pt" +[ -e "$DEFAULT_WEIGHTS" ] || DEFAULT_WEIGHTS="$REPO/runs/detect/film/weights/best.pt" + +export YOLO_WEIGHTS="${YOLO_WEIGHTS:-$DEFAULT_WEIGHTS}" +export YOLO_DEVICE="${YOLO_DEVICE:-cuda}" +export YOLO_CONF="${YOLO_CONF:-0.25}" +export YOLO_IMGSZ="${YOLO_IMGSZ:-640}" + +# 必须与运行 Label Studio 时的值一致，否则按 ?d=... 还原不出图片路径 +export LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT="${LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT:-/home/deep/labeling}" + +# 点 Start Training 无条件触发；攒够 RETRAIN_AFTER 条新标注也会自动触发（0=只手动） +export RETRAIN_AFTER="${RETRAIN_AFTER:-0}" +export RETRAIN_COMMAND="${RETRAIN_COMMAND:-$REPO/scripts/retrain_from_label_studio.sh}" +export RETRAIN_LOG="${RETRAIN_LOG:-/tmp/retrain.log}" + +PORT="${LS_ML_PORT:-9090}" +LOG="${ML_BACKEND_LOG:-/tmp/mlb.log}" +PYTHON="${ML_BACKEND_PYTHON:-$REPO/.venv-ml/bin/python}" + +if [ ! -x "$PYTHON" ]; then + echo "找不到 $PYTHON。先建独立环境（见 ml_backend/README.md）：" >&2 + echo " uv venv .venv-ml --python 3.12" >&2 + exit 1 +fi +if [ ! -e "$YOLO_WEIGHTS" ]; then + echo "找不到权重 $YOLO_WEIGHTS" >&2 + exit 1 +fi + +# 模式里不要出现 _wsgi.py 的完整写法：ssh 单行执行时会匹配到命令自身，把会话一起杀掉 +pkill -f "wsgi" 2>/dev/null || true +sleep 1 + +cd "$REPO/ml_backend" +setsid nohup "$PYTHON" ./_wsgi.py -p "$PORT" > "$LOG" 2>&1 < /dev/null & +echo "ML backend 已启动 pid=$!" +echo " 权重 : $YOLO_WEIGHTS" +echo " 重训命令 : $RETRAIN_COMMAND (RETRAIN_AFTER=$RETRAIN_AFTER)" +echo " 服务日志 : $LOG" +echo " 重训日志 : $RETRAIN_LOG" +echo " 网页端 URL: http://localhost:$PORT" diff --git a/tests/test_ml_backend_convert.py b/tests/test_ml_backend_convert.py index 15d57cc..a0ae64c 100644 --- a/tests/test_ml_backend_convert.py +++ b/tests/test_ml_backend_convert.py @@ -69,3 +69,37 @@ def test_score_is_kept_per_box(): [("blackdot", [0.1, 0.1, 0.2, 0.2], 0.42)], 64, 64, FROM_NAME, TO_NAME, "v1" ) assert prediction["result"][0]["score"] == pytest.approx(0.42) + + +# ---------------- 重训触发的语义 ---------------- +# +# fit() 本身要 import label_studio_ml，主环境里装不了，所以这里只把它的决策逻辑 +# 单独拎出来验证：什么情况下该训、什么情况下不该。规则见 ml_backend/model.py。 + + +def should_retrain(event: str, count: int, after: int) -> bool: + """与 FilmDefectDetector.fit 中的判定保持一致。""" + if event == "START_TRAINING": + return True + return bool(after and count >= after) + + +@pytest.mark.parametrize( + "event,count,after,expected", + [ + ("START_TRAINING", 0, 0, True), # 人点的按钮，无条件训 + ("START_TRAINING", 3, 200, True), # 没攒够也训 + ("ANNOTATION_CREATED", 199, 200, False), + ("ANNOTATION_CREATED", 200, 200, True), + ("ANNOTATION_CREATED", 999, 0, False), # after=0 表示只手动触发 + ], +) +def test_retrain_trigger_rules(event, count, after, expected): + assert should_retrain(event, count, after) is expected + + +def test_only_new_annotations_count_toward_the_threshold(): + """UPDATED / DELETED 是对已有标注的修改，不该把阈值提前顶到。""" + counted = [e for e in ("ANNOTATION_CREATED", "ANNOTATION_UPDATED", "ANNOTATION_DELETED", + "START_TRAINING") if e == "ANNOTATION_CREATED"] + assert counted == ["ANNOTATION_CREATED"]