| """权重瘦身: 只保留推理需要的部分。 |
| |
| 推理的加载逻辑是「先 load model_state_dict, 再用 ema_state_dict.shadow 逐参数覆盖」, |
| 所以把 EMA 影子权重合并进 model_state_dict 后只存一份, 推理结果逐位等价。 |
| |
| 丢弃: optimizer_state_dict(最大头, AdamW 的 m/v 约 2x 参数量)、scheduler_state_dict、训练元信息。 |
| |
| 用法: python tools_strip_ckpt.py <输入.pth> <输出.pth> [--no-ema] |
| """ |
| import os, sys |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
|
|
| import argparse, os, torch |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("src"); ap.add_argument("dst") |
| ap.add_argument("--no-ema", action="store_true", help="不合并 EMA, 保留原始训练权重") |
| a = ap.parse_args() |
|
|
| ck = torch.load(a.src, map_location="cpu", weights_only=False) |
| print(f"输入 {a.src} ({os.path.getsize(a.src)/1e9:.2f} GB)") |
| print(" 顶层 key:", list(ck.keys())) |
| msd = ck["model_state_dict"] |
| shadow = (ck.get("ema_state_dict") or {}).get("shadow") if isinstance(ck.get("ema_state_dict"), dict) else None |
|
|
| merged = 0 |
| if shadow and not a.no_ema: |
| for k in msd: |
| if k in shadow: |
| msd[k] = shadow[k].to(msd[k].dtype); merged += 1 |
| print(f" 已合并 EMA 影子权重: {merged}/{len(msd)} 个张量") |
| elif a.no_ema: |
| print(" --no-ema: 保留原始训练权重") |
| else: |
| print(" ckpt 内无 EMA, 直接使用 model_state_dict") |
|
|
| out = {"model_state_dict": msd} |
| torch.save(out, a.dst) |
| print(f"输出 {a.dst} ({os.path.getsize(a.dst)/1e9:.2f} GB, " |
| f"省 {100*(1-os.path.getsize(a.dst)/os.path.getsize(a.src)):.0f}%)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
