File size: 1,839 Bytes
90ca518 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """权重瘦身: 只保留推理需要的部分。
推理的加载逻辑是「先 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()
|