312 lines
12 KiB
Python
312 lines
12 KiB
Python
#!/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))
|
||
ttype = fb.field(t, 1)
|
||
return {"name": fb.string(fb.field(t, 3)) or "?",
|
||
"shape": shape,
|
||
"type": TENSOR_TYPE.get(fb.d[ttype], fb.d[ttype]) if ttype else "?"}
|
||
|
||
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()
|