| """Average the weights of two epoch checkpoints (the z2 model is the 50/50 average of epochs 0 and 1). | |
| python average_checkpoints.py runs/z2/epoch0.pt runs/z2/epoch1.pt --out runs/z2/avg.pt | |
| """ | |
| import argparse | |
| import torch | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("a") | |
| ap.add_argument("b") | |
| ap.add_argument("--out", required=True) | |
| args = ap.parse_args() | |
| a = torch.load(args.a, map_location="cpu", weights_only=False) | |
| b = torch.load(args.b, map_location="cpu", weights_only=False) | |
| assert a["config"] == b["config"] and a["model"].keys() == b["model"].keys() | |
| avg = {k: (0.5 * a["model"][k].float() + 0.5 * b["model"][k].float()).to(a["model"][k].dtype) | |
| if a["model"][k].is_floating_point() else a["model"][k] for k in a["model"]} | |
| torch.save({"model": avg, "config": a["config"]}, args.out) | |