1

server/training/train_server.py

@@ -34,6 +34,10 @@ import sys
 import traceback
 from pathlib import Path
 
+# MPS（Apple Silicon）训练：默认水位上限把可用统一内存压在 ~2/3 且缓存不归还，16GB 机器易 OOM；
+# 须在 torch 加载前取消上限，配合每 epoch 手动 empty_cache 归还（Go 侧 device=mps 时生效）
+os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
+
 # 抗尺度漂移（与人工训练基线一致）：随机缩放输入 0.5~1.5x
 MULTI_SCALE = 0.5
 # 早停耐心（连续 N 轮无提升即停）
@@ -241,6 +245,7 @@ def main():
     epochs = int(task.get("epochs") or 150)
     batch = int(task.get("batch") or 16)
     device = task.get("device") or "0"
+    is_mps = str(device).strip().lower() == "mps"
     data = os.path.join(task["yolo"], "dataset.yaml")
     names = read_names(task["yolo"])
 
@@ -248,13 +253,49 @@ def main():
         from ultralytics import YOLO
         model = YOLO(task.get("model") or "yolov8s.pt")
         register_callback(model)
-        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,
-        )
+        if is_mps:
+            def _mps_cache(_trainer):
+                # MPS 分配器缓存不自动归还，每 epoch 清一次防碎片化 OOM
+                try:
+                    import torch
+                    torch.mps.empty_cache()
+                except Exception:
+                    pass
+            model.add_callback("on_fit_epoch_end", _mps_cache)
+
+        # batch OOM 自适应：统一内存/显存不足时 batch 减半重试（最低 1），速度换可跑
+        cur_batch = batch
+        while True:
+            try:
+                model.train(
+                    data=data, imgsz=imgsz, epochs=epochs,
+                    patience=PATIENCE, batch=cur_batch, device=device,
+                    # macOS dataloader worker 为 spawn 子进程、各占一份内存，MPS 直接不用
+                    workers=0 if is_mps else 4,
+                    # project 必须绝对路径，相对路径会被拼到默认 runs/detect 下造成双层嵌套
+                    project=task_path("project"), name="train",
+                    exist_ok=True, plots=True,
+                    # 多尺度随机放大输入（最高 1.5x imgsz）是 MPS 内存尖峰来源，关闭
+                    multi_scale=0 if is_mps else MULTI_SCALE,
+                )
+                break
+            except Exception as e:
+                low = f"{e}".lower()
+                oom = ("out of memory" in low
+                       or "outofmemory" in type(e).__name__.lower()
+                       or ("mps" in low and "memory" in low))
+                if not oom or cur_batch <= 1:
+                    raise
+                cur_batch = max(1, cur_batch // 2)
+                print(f"[train] 内存不足（OOM），batch 降为 {cur_batch} 重试", file=sys.stderr)
+                try:
+                    import torch
+                    if torch.backends.mps.is_available():
+                        torch.mps.empty_cache()
+                    elif torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                except Exception:
+                    pass
         trainer = getattr(model, "trainer", None)
         if trainer is None:
             raise RuntimeError("训练完成但无法获取 trainer（save_dir 未知）")