Spaces:
Sleeping
Sleeping
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.") | |