93 lines
4.6 KiB
Python
93 lines
4.6 KiB
Python
"""断点存档与"体面停止"（2026-10-18 第五轮补充：用户要求长训练意外中断后能接着跑，不从头来）。
|
||
|
||
- atomic_save：先写 <名>.tmp 再 os.replace —— 断电/被杀发生在写盘中途时，旧的 last.pt 仍完整可用。
|
||
- Saver.due()：按墙钟每 every_min 分钟存一次 last.pt（不再只在每 2000 步验证时存，断电最多丢 10 分钟）。
|
||
- Saver.stop()：收到 SIGTERM/SIGINT，或看到停止标记文件 ml/runs/STOP 时返回 True；
|
||
训练循环据此先存 last.pt 再以 143 退出。Windows 上 taskkill /F 是硬杀、收不到信号，
|
||
所以控制台「停止」先写 STOP 文件、等进程自己存档退出，超时才硬杀（见 monitor.Job.stop）。
|
||
"""
|
||
import os, random, signal, time
|
||
import torch
|
||
|
||
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
STOP = os.path.join(ROOT, 'ml', 'runs', 'STOP')
|
||
|
||
|
||
def atomic_save(obj, path):
|
||
tmp = path + '.tmp'
|
||
torch.save(obj, tmp)
|
||
os.replace(tmp, path)
|
||
|
||
|
||
class Saver:
|
||
def __init__(self, every_min=10.0):
|
||
self.every = every_min * 60; self.last = time.time(); self._sig = False; self._chk = 0.0
|
||
for s in ('SIGTERM', 'SIGINT', 'SIGBREAK', 'SIGHUP'):
|
||
if hasattr(signal, s):
|
||
try: signal.signal(getattr(signal, s), self._on)
|
||
except (ValueError, OSError): pass # 非主线程等情况：只靠 STOP 文件
|
||
|
||
def _on(self, *_): self._sig = True
|
||
|
||
def stop(self):
|
||
if self._sig: return True
|
||
now = time.time()
|
||
if now - self._chk > 2: # 文件每 2 s 查一次，不拖慢训练步
|
||
self._chk = now; self._sig = os.path.exists(STOP)
|
||
return self._sig
|
||
|
||
def due(self):
|
||
if time.time() - self.last >= self.every:
|
||
self.last = time.time(); return True
|
||
return False
|
||
|
||
def mark(self): self.last = time.time()
|
||
|
||
|
||
def rng_state():
|
||
"""2026-10-18 第七轮（遗留：续跑后数值轨迹与不断点不一致）：随机数状态一起存档。
|
||
训练取 batch 用 torch.randint（全局生成器，CUDA 上取 CUDA 生成器）、噪声用 torch.randn、APG 选题用 random.choices，
|
||
三者状态都恢复后，续跑的 batch 顺序与噪声与一口气跑完一致（同一设备、同一版本；cuDNN 非确定算子另计）。"""
|
||
d = {'torch': torch.get_rng_state(), 'py': random.getstate()}
|
||
if torch.cuda.is_available(): d['cuda'] = torch.cuda.get_rng_state_all()
|
||
return d # 不存 numpy 状态：训练不用 np.random，且 ndarray 过不了 torch.load 的 weights_only 默认
|
||
|
||
|
||
def set_rng_state(d):
|
||
if not d: return False
|
||
torch.set_rng_state(d['torch'].cpu() if hasattr(d['torch'], 'cpu') else d['torch']); random.setstate(d['py'])
|
||
if 'cuda' in d and torch.cuda.is_available():
|
||
try: torch.cuda.set_rng_state_all([t.cpu() for t in d['cuda']])
|
||
except Exception: pass
|
||
return True
|
||
|
||
|
||
def load_policy(path, device, prefer_ema=True):
|
||
"""载入策略（2026-10-19，审查 review_model_train P0-2）：优先取 EMA 权重 ck['ema']（train.py 起存），没有就用 ck['model']（旧 checkpoint、APG 产物）。
|
||
model.act_scale = ck['act_scale']（动作标签归一化倍数，见 ml/prep/build_dataset.cjs ACT_SCALE）；旧 checkpoint 缺省 1.0 = 原口径。
|
||
闭环仿真 / DAgger / APG 起点 / 挖错 / 门禁 / 验收 / 导出都经这里载入，保证同一份权重、同一个动作尺度。返回 (model.eval(), ck)。"""
|
||
from ml.model.policy import build
|
||
ck = torch.load(path, map_location=device)
|
||
m = build(ck['cfg']).to(device)
|
||
sd = ck.get('ema') if prefer_ema and ck.get('ema') is not None else ck['model']
|
||
m.load_state_dict(sd)
|
||
m.act_scale = float(ck.get('act_scale', 1.0))
|
||
return m.eval(), ck
|
||
|
||
|
||
def pick_start(run_dir):
|
||
"""DAgger 起点（2026-10-19，审查 P0-2）：best.pt 的闭环 success > 0 才用 best.pt；否则用 last.pt
|
||
（success 全 0 时 best.pt 只由开环误差决定，可能落在 p_neq 预热期，一步映射没训过）。载入时取 EMA。"""
|
||
best, last = os.path.join(run_dir, 'best.pt'), os.path.join(run_dir, 'last.pt')
|
||
if os.path.exists(best):
|
||
ev = torch.load(best, map_location='cpu').get('eval') or {}
|
||
if (ev.get('success') or 0) > 0 or not os.path.exists(last): return best
|
||
return last if os.path.exists(last) else best
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# python -m ml.common.ckpt pick <运行目录> → 打印 DAgger 起点 checkpoint 路径
|
||
import sys
|
||
if len(sys.argv) == 3 and sys.argv[1] == 'pick': print(pick_start(sys.argv[2]))
|
||
else: raise SystemExit('usage: python -m ml.common.ckpt pick <run_dir>')
|