| import torch |
| from collections import OrderedDict |
|
|
|
|
| def extract(ckpt): |
| model = ckpt["model"] |
| opt = OrderedDict() |
| opt["weight"] = {key: value for key, value in model.items() if "enc_q" not in key} |
| return opt |
|
|
|
|
| def model_fusion(model_name, pth_path_1, pth_path_2): |
| ckpt1 = torch.load(pth_path_1, map_location="cpu") |
| ckpt2 = torch.load(pth_path_2, map_location="cpu") |
| if "model" in ckpt1: |
| ckpt1 = extract(ckpt1) |
| else: |
| ckpt1 = ckpt1["weight"] |
| if "model" in ckpt2: |
| ckpt2 = extract(ckpt2) |
| else: |
| ckpt2 = ckpt2["weight"] |
| if sorted(ckpt1.keys()) != sorted(ckpt2.keys()): |
| return "Fail to merge the models. The model architectures are not the same." |
| opt = OrderedDict( |
| weight={ |
| key: 1 * value.float() + (1 - 1) * ckpt2[key].float() |
| for key, value in ckpt1.items() |
| } |
| ) |
| opt["info"] = f"Model fusion of {pth_path_1} and {pth_path_2}" |
| torch.save(opt, f"logs/{model_name}.pth") |
| print(f"Model fusion of {pth_path_1} and {pth_path_2} is done.") |
|
|
