训练体系整合与标注单阶段化
- 标注：AI 预标注直写 labels_json（去候选确认两阶段）；重叠去重（minIoU）；全量标注按钮 - 训练：脚本迁移入 server/training/（Go 化 prepare_yolo/analyze_rfdetr，保留 train_server.py）；tflite 产物自检并入训练流程（check_tflite） - 数据目录/权重不进 git；.gitignore 迁移至仓库根
This commit is contained in:
@@ -0,0 +1,313 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""训练任务脚本（Go 后端 --task-json 驱动，产物契约见 server/common/training_runner.go）。
|
||||
|
||||
用法：
|
||||
python train_server.py --task-json tasks/<taskId>.json
|
||||
|
||||
任务参数（Go 侧写入，字段相对训练机 workdir）：
|
||||
workdir 训练机工作目录（脚本 / yolov8n.pt / venv 所在），启动即 chdir
|
||||
yolo 训练集目录（含 dataset.yaml），相对 workdir
|
||||
imgsz 训练/导出分辨率（默认 704，与端侧推理对齐）
|
||||
epochs 训练轮数
|
||||
batch 批大小
|
||||
device GPU 编号或 cpu
|
||||
project 训练输出目录（相对 workdir，末尾自动拼 name）
|
||||
log_file 每 epoch 一行 JSON 的进度日志（相对 workdir）
|
||||
result_file 结束结果 JSON（相对 workdir）
|
||||
artifact_zip 产物打包（best.pt + results.csv + 曲线，相对 workdir）
|
||||
|
||||
产物契约：
|
||||
log_file {"epoch":1,"total":150,"metrics":{"metrics/mAP50(B)":0.87,...}}
|
||||
result_file {"metrics":{...},"names":["pheasant","suspect"],"best_tflite":"runs/tasks/<id>/weights/best.tflite",
|
||||
"tflite_check":{"ok":true,"reason":"","inputs":[...],"outputs":[...]}}
|
||||
result_file 存在 = 训练完成；异常时写 {"error":"..."}，Go 侧据以置失败并展示原因。
|
||||
tflite_check.ok=false（产物 shape 异常）时 Go 侧置训练失败并带出 reason。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
import traceback
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
# 抗尺度漂移（与人工训练基线一致）：随机缩放输入 0.5~1.5x
|
||||
MULTI_SCALE = 0.5
|
||||
# 早停耐心（连续 N 轮无提升即停）
|
||||
PATIENCE = 30
|
||||
|
||||
|
||||
def on_fit_epoch_end(trainer):
|
||||
"""ultralytics 回调：每 epoch 结束写一行进度 JSON（行缓冲，进程被杀不丢行）。"""
|
||||
metrics = {}
|
||||
for k, v in (trainer.metrics or {}).items():
|
||||
if isinstance(v, (int, float)) and math.isfinite(v):
|
||||
metrics[k] = round(float(v), 5)
|
||||
if _LOG is None:
|
||||
return
|
||||
total = getattr(trainer, "epochs", None) or getattr(trainer, "total_epochs", 0) or 0
|
||||
_LOG.write(json.dumps({
|
||||
"epoch": trainer.epoch + 1,
|
||||
"total": total,
|
||||
"metrics": metrics,
|
||||
}, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def register_callback():
|
||||
from ultralytics.utils.callbacks import callbacks
|
||||
callbacks["on_fit_epoch_end"].append(on_fit_epoch_end)
|
||||
|
||||
|
||||
def write_result(result_file, payload):
|
||||
Path(result_file).parent.mkdir(parents=True, exist_ok=True)
|
||||
# tmp + rename 原子写，避免 Go 侧读到半截文件
|
||||
tmp = result_file + ".tmp"
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, ensure_ascii=False, indent=2)
|
||||
os.replace(tmp, result_file)
|
||||
|
||||
|
||||
def read_names(yolo_dir):
|
||||
"""从 data.yaml 读类别名（有序列表，ultralytics 依赖 pyyaml）"""
|
||||
try:
|
||||
import yaml
|
||||
with open(Path(yolo_dir) / "dataset.yaml", encoding="utf-8") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
names = cfg.get("names") or {}
|
||||
if isinstance(names, dict):
|
||||
return [names[k] for k in sorted(names, key=lambda k: int(k))]
|
||||
return list(names)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def find_best_tflite(save_dir):
|
||||
"""定位导出的 tflite（ultralytics 各版本产物位置不一）：
|
||||
优先 fp32 best.tflite（与人工基线一致），再 float32/int8，最后兜底全局搜。"""
|
||||
weights = save_dir / "weights"
|
||||
for name in ("best.tflite", "best_float32.tflite", "best_int8.tflite"):
|
||||
p = weights / name
|
||||
if p.exists():
|
||||
return p
|
||||
candidates = sorted(weights.rglob("*.tflite"))
|
||||
if candidates:
|
||||
return candidates[0]
|
||||
return None
|
||||
|
||||
|
||||
# TFLite flatbuffer 张量类型（tensorflow/lite/schema/schema.fbs，字段顺序敏感）
|
||||
TENSOR_TYPE = {0: "FLOAT32", 1: "FLOAT16", 2: "INT32", 3: "UINT8", 4: "INT64",
|
||||
5: "STRING", 6: "BOOL", 7: "INT16", 8: "COMPLEX64", 9: "INT8",
|
||||
10: "FLOAT64", 11: "COMPLEX128", 12: "UINT64", 13: "RESOURCE",
|
||||
14: "VARIANT", 15: "UINT32", 16: "UINT16", 17: "INT4",
|
||||
18: "BFLOAT16", 19: "FLOAT8_E4M3FN", 20: "FLOAT8_E4M3FNUZ",
|
||||
21: "FLOAT8_E5M2", 22: "FLOAT8_E5M2FNUZ"}
|
||||
|
||||
|
||||
class _FB:
|
||||
"""极简 flatbuffer 读取器（仅 Model/SubGraph/Tensor 表所需字段）"""
|
||||
|
||||
def __init__(self, data):
|
||||
self.d = data
|
||||
|
||||
def u32(self, pos): return struct.unpack_from("<I", self.d, pos)[0]
|
||||
def i32(self, pos): return struct.unpack_from("<i", self.d, pos)[0]
|
||||
def u16(self, pos): return struct.unpack_from("<H", self.d, pos)[0]
|
||||
|
||||
def field(self, t, i):
|
||||
vto = self.i32(t)
|
||||
vt = t - vto
|
||||
vs = self.u16(vt)
|
||||
off = vt + 4 + 2 * i
|
||||
if off + 2 > vt + vs:
|
||||
return None
|
||||
f = self.u16(off)
|
||||
return None if f == 0 else t + f
|
||||
|
||||
def deref(self, pos):
|
||||
return pos + self.u32(pos) if pos is not None else None
|
||||
|
||||
def vec(self, pos):
|
||||
pos = self.deref(pos)
|
||||
if pos is None:
|
||||
return None
|
||||
return pos + 4, self.u32(pos)
|
||||
|
||||
def vec_table(self, pos):
|
||||
start, n = self.vec(pos)
|
||||
return [start + i * 4 + self.u32(start + i * 4) for i in range(n)]
|
||||
|
||||
def string(self, pos):
|
||||
pos = self.deref(pos)
|
||||
if pos is None:
|
||||
return None
|
||||
n = self.u32(pos)
|
||||
return self.d[pos + 4:pos + 4 + n].decode("utf-8", "replace")
|
||||
|
||||
def int_vec(self, pos):
|
||||
start, n = self.vec(pos)
|
||||
return [self.i32(start + 4 * i) for i in range(n)]
|
||||
|
||||
|
||||
def check_tflite(path, imgsz):
|
||||
"""TFLite 产物自检（原 inspect_tflite.py，2026-08-26 并入）：
|
||||
输入恰 1 张且 4 维、元素总数 == imgsz²×3（NCHW/NHWC 皆可）、输出 ≥ 1 张且 batch 维 = 1。
|
||||
返回 {"ok","reason","inputs":[{"name","shape","type"}],"outputs":[...]}。"""
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
fb = _FB(f.read())
|
||||
root = fb.u32(0) # Model 表
|
||||
subgraphs = fb.vec_table(fb.field(root, 2)) # Model.subgraphs [2]
|
||||
sg = subgraphs[0]
|
||||
tensors = fb.vec_table(fb.field(sg, 0)) # SubGraph.tensors [0]
|
||||
inputs = fb.int_vec(fb.field(sg, 1)) # SubGraph.inputs [1]
|
||||
outputs = fb.int_vec(fb.field(sg, 2)) # SubGraph.outputs [2]
|
||||
|
||||
def describe(i):
|
||||
t = tensors[i]
|
||||
shape = fb.int_vec(fb.field(t, 0))
|
||||
# type 有默认值 FLOAT32(=0)：字段等于默认值时 flatbuffer 省略（vtable 偏移 0）
|
||||
ttype_f = fb.field(t, 1)
|
||||
ttype = fb.d[ttype_f] if ttype_f else 0
|
||||
return {"name": fb.string(fb.field(t, 3)) or "?",
|
||||
"shape": shape,
|
||||
"type": TENSOR_TYPE.get(ttype, ttype)}
|
||||
|
||||
in_desc = [describe(i) for i in inputs]
|
||||
out_desc = [describe(i) for i in outputs]
|
||||
|
||||
problems = []
|
||||
if len(in_desc) != 1:
|
||||
problems.append(f"输入张量 {len(in_desc)} 个（期望 1）")
|
||||
else:
|
||||
shape = in_desc[0]["shape"]
|
||||
n = 1
|
||||
for d in shape:
|
||||
n *= d
|
||||
if len(shape) != 4:
|
||||
problems.append(f"输入 shape {shape} 非 4 维")
|
||||
elif n != imgsz * imgsz * 3:
|
||||
problems.append(f"输入元素数 {n} ≠ imgsz²×3={imgsz * imgsz * 3}（shape {shape}，imgsz={imgsz}）")
|
||||
if not out_desc:
|
||||
problems.append("无输出张量")
|
||||
else:
|
||||
for o in out_desc:
|
||||
if o["shape"] and o["shape"][0] != 1:
|
||||
problems.append(f"输出 batch 维 ≠ 1: {o['shape']}")
|
||||
|
||||
return {"ok": not problems, "reason": "；".join(problems),
|
||||
"inputs": in_desc, "outputs": out_desc}
|
||||
except Exception as e:
|
||||
return {"ok": False, "reason": f"解析失败: {e}", "inputs": [], "outputs": []}
|
||||
|
||||
|
||||
def build_artifact_zip(zip_path, save_dir):
|
||||
Path(zip_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for p in (save_dir / "weights" / "best.pt", save_dir / "results.csv"):
|
||||
if p.exists():
|
||||
zf.write(p, arcname=p.name)
|
||||
# 曲线/混淆矩阵等绘图产物在 save_dir 根目录
|
||||
for p in sorted(save_dir.glob("*.png")) + sorted(save_dir.glob("*.jpg")):
|
||||
zf.write(p, arcname=p.name)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="observer 训练任务脚本（Go --task-json 驱动）")
|
||||
ap.add_argument("--task-json", required=True, help="任务参数 JSON 文件路径")
|
||||
args = ap.parse_args()
|
||||
|
||||
with open(args.task_json, encoding="utf-8") as f:
|
||||
task = json.load(f)
|
||||
|
||||
global _LOG
|
||||
base = os.path.abspath(task["workdir"])
|
||||
|
||||
def task_path(key):
|
||||
"""任务参数里的路径全部相对 workdir，这里绝对化（chdir 失败也能定位）"""
|
||||
p = task[key]
|
||||
return p if os.path.isabs(p) else os.path.join(base, p)
|
||||
|
||||
try:
|
||||
os.chdir(base)
|
||||
except Exception:
|
||||
pass # 目录不可用则后续训练必然失败，错误路径仍尽力写 result.json
|
||||
|
||||
log_file = task_path("log_file")
|
||||
result_file = task_path("result_file")
|
||||
artifact_zip = task_path("artifact_zip")
|
||||
global _LOG
|
||||
_LOG = None
|
||||
try:
|
||||
Path(log_file).parent.mkdir(parents=True, exist_ok=True)
|
||||
_LOG = open(log_file, "a", encoding="utf-8", buffering=1)
|
||||
except Exception:
|
||||
pass # 日志不可写（如 workdir 缺失）不阻断错误路径，traceback 仍走 stderr
|
||||
|
||||
imgsz = int(task.get("imgsz") or 704)
|
||||
epochs = int(task.get("epochs") or 150)
|
||||
batch = int(task.get("batch") or 16)
|
||||
device = task.get("device") or "0"
|
||||
data = os.path.join(task["yolo"], "dataset.yaml")
|
||||
names = read_names(task["yolo"])
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
register_callback()
|
||||
model = YOLO("yolov8n.pt")
|
||||
model.train(
|
||||
data=data, imgsz=imgsz, epochs=epochs,
|
||||
patience=PATIENCE, batch=batch, device=device, workers=4,
|
||||
# project 必须绝对路径，相对路径会被拼到默认 runs/detect 下造成双层嵌套
|
||||
project=task_path("project"), name="train",
|
||||
exist_ok=True, plots=True, multi_scale=MULTI_SCALE,
|
||||
)
|
||||
trainer = getattr(model, "trainer", None)
|
||||
if trainer is None:
|
||||
raise RuntimeError("训练完成但无法获取 trainer（save_dir 未知）")
|
||||
save_dir = Path(trainer.save_dir)
|
||||
|
||||
best_tflite = find_best_tflite(save_dir)
|
||||
if best_tflite is None:
|
||||
model.export(format="tflite", imgsz=imgsz)
|
||||
best_tflite = find_best_tflite(save_dir)
|
||||
if best_tflite is None:
|
||||
raise RuntimeError("tflite 导出失败：weights 目录下未找到任何 .tflite 产物")
|
||||
|
||||
# tflite 产物自检：shape 异常也写进 result.json（ok=false），Go 侧据此置失败
|
||||
tflite_check = check_tflite(str(best_tflite), imgsz)
|
||||
if not tflite_check["ok"]:
|
||||
print(f"[tflite-check] FAIL: {tflite_check['reason']}", file=sys.stderr)
|
||||
|
||||
# 结果指标取 best 轮（发布的是 best.pt），缺省回退末轮
|
||||
metrics = getattr(trainer, "best_metrics", None) or trainer.metrics
|
||||
clean_metrics = {}
|
||||
for k, v in (metrics or {}).items():
|
||||
if isinstance(v, (int, float)) and math.isfinite(v):
|
||||
clean_metrics[k] = round(float(v), 5)
|
||||
|
||||
build_artifact_zip(artifact_zip, save_dir)
|
||||
# best_tflite 相对 workdir，Go 侧按此路径拉取
|
||||
write_result(result_file, {
|
||||
"metrics": clean_metrics,
|
||||
"names": names,
|
||||
"best_tflite": os.path.relpath(best_tflite, base),
|
||||
"tflite_check": tflite_check,
|
||||
})
|
||||
except Exception:
|
||||
err = traceback.format_exc()
|
||||
if _LOG is not None:
|
||||
_LOG.write("\n" + err + "\n")
|
||||
try:
|
||||
write_result(result_file, {"error": err[-2000:]})
|
||||
except Exception:
|
||||
pass
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user