File size: 1,651 Bytes
c668e80 |
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 53 54 55 56 57 58 59 60 61 |
#!/usr/bin/env python
import argparse
import torch
def average_models(model_files, fp32=False):
vocab = None
opt = None
avg_model = None
avg_generator = None
for i, model_file in enumerate(model_files):
m = torch.load(model_file, map_location="cpu")
model_weights = m["model"]
generator_weights = m["generator"]
if fp32:
for k, v in model_weights.items():
model_weights[k] = v.float()
for k, v in generator_weights.items():
generator_weights[k] = v.float()
if i == 0:
vocab, opt = m["vocab"], m["opt"]
avg_model = model_weights
avg_generator = generator_weights
else:
for k, v in avg_model.items():
avg_model[k].mul_(i).add_(model_weights[k]).div_(i + 1)
for k, v in avg_generator.items():
avg_generator[k].mul_(i).add_(generator_weights[k]).div_(i + 1)
final = {
"vocab": vocab,
"opt": opt,
"optim": None,
"generator": avg_generator,
"model": avg_model,
}
return final
def main():
parser = argparse.ArgumentParser(description="")
parser.add_argument(
"-models", "-m", nargs="+", required=True, help="List of models"
)
parser.add_argument("-output", "-o", required=True, help="Output file")
parser.add_argument(
"-fp32", "-f", action="store_true", help="Cast params to float32"
)
opt = parser.parse_args()
final = average_models(opt.models, opt.fp32)
torch.save(final, opt.output)
if __name__ == "__main__":
main()
|