|
|
|
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() |
|
|
