| |
| """SWA: son N intermediate ckpt'in ağırlık ortalaması. |
| |
| Usage: python swa_average.py runs/detect/.../yolo_fold0/weights/ [N=5] |
| """ |
| import sys, glob, copy |
| from pathlib import Path |
| import torch |
|
|
| def main(): |
| if len(sys.argv) < 2: |
| print("usage: swa_average.py <weights_dir> [N=5]"); sys.exit(1) |
| wdir = Path(sys.argv[1]) |
| n = int(sys.argv[2]) if len(sys.argv) > 2 else 5 |
|
|
| |
| ckpts = sorted(wdir.glob('epoch*.pt')) |
| if len(ckpts) < n: |
| print(f"only {len(ckpts)} epoch ckpts; using all + last.pt"); |
| last = wdir / 'last.pt' |
| if last.exists() and last not in ckpts: |
| ckpts.append(last) |
| ckpts = ckpts[-n:] |
| print(f"Averaging {len(ckpts)}: {[c.name for c in ckpts]}") |
|
|
| avg_state = None; avg_meta = None |
| for i, c in enumerate(ckpts): |
| d = torch.load(c, map_location='cpu') |
| sd = d['model'].state_dict() if hasattr(d['model'], 'state_dict') else d['model'] |
| if i == 0: |
| avg_state = {k: v.float().clone() for k, v in sd.items()} |
| avg_meta = d |
| else: |
| for k in avg_state: |
| if k in sd and avg_state[k].shape == sd[k].shape: |
| avg_state[k] += sd[k].float() |
| for k in avg_state: |
| avg_state[k] /= len(ckpts) |
|
|
| |
| if hasattr(avg_meta['model'], 'load_state_dict'): |
| avg_meta['model'].load_state_dict({k: v.to(next(avg_meta['model'].parameters()).dtype) for k,v in avg_state.items()}) |
| else: |
| avg_meta['model'] = avg_state |
|
|
| out = wdir / 'swa.pt' |
| torch.save(avg_meta, out) |
| print(f"SWA -> {out}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
