| """scripts/train.py を包む。Trainer が「DiT 以外を凍結」した直後に、環境変数 FLEX_FREEZE_REGEX に |
| 合うパラメータを追加で凍結する（F2 = 動画専門家 mot.mixtures.video を凍結し行動側だけ学習）。 |
| upstream は編集しない。学習可能パラメータ数を印字して、凍結が実際に効いたことを目で確かめる。""" |
| import os, re, sys, runpy |
| sys.path.insert(0, "/workspace/flexpi_repo/src") |
| import flexpi.trainer as T |
|
|
| _orig = T.Wan22Trainer.__dict__["_apply_dit_only_train_mode"].__func__ |
|
|
| def _patched(model): |
| _orig(model) |
| pat = os.environ.get("FLEX_FREEZE_REGEX", "") |
| frozen = 0 |
| if pat: |
| rx = re.compile(pat) |
| for n, p in model.named_parameters(): |
| if p.requires_grad and rx.search(n): |
| p.requires_grad_(False); frozen += p.numel() |
| tr = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| tot = sum(p.numel() for p in model.parameters()) |
| heads = {} |
| for n, p in model.named_parameters(): |
| if p.requires_grad: |
| k = ".".join(n.split(".")[:3]); heads[k] = heads.get(k, 0) + p.numel() |
| top = sorted(heads.items(), key=lambda x: -x[1])[:8] |
| print(f"[wrap] 学習可能 {tr/1e9:.3f}B / 全体 {tot/1e9:.3f}B regex 凍結 {frozen/1e9:.3f}B pattern={pat!r}", flush=True) |
| print("[wrap] 学習可能の内訳（上位）: " + " ".join(f"{k}={v/1e6:.0f}M" for k, v in top), flush=True) |
|
|
| T.Wan22Trainer._apply_dit_only_train_mode = staticmethod(_patched) |
| os.chdir("/workspace/flexpi_repo") |
| sys.argv = ["scripts/train.py"] + sys.argv[1:] |
| runpy.run_path("scripts/train.py", run_name="__main__") |
|
|
