File size: 1,727 Bytes
f211247 | 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 48 49 50 51 52 | #!/usr/bin/env python3
"""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
# Collect epoch ckpts (epoch{N}.pt or last.pt)
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)
# Replace model state, save
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()
|