| 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.") | |