| New file |
| | |
|---|
| | | #!/usr/bin/env python3 |
|---|
| | | |
|---|
| | | import argparse |
|---|
| | | |
|---|
| | | import torch |
|---|
| | | |
|---|
| | | |
|---|
| | | def average_model(input_files, output_file): |
|---|
| | | output_model = {} |
|---|
| | | for ckpt_path in input_files: |
|---|
| | | model_params = torch.load(ckpt_path, map_location="cpu") |
|---|
| | | for key, value in model_params.items(): |
|---|
| | | if key not in output_model: |
|---|
| | | output_model[key] = value |
|---|
| | | else: |
|---|
| | | output_model[key] += value |
|---|
| | | for key in output_model.keys(): |
|---|
| | | output_model[key] /= len(input_files) |
|---|
| | | torch.save(output_model, output_file) |
|---|
| | | |
|---|
| | | |
|---|
| | | if __name__ == '__main__': |
|---|
| | | parser = argparse.ArgumentParser() |
|---|
| | | parser.add_argument("output_file") |
|---|
| | | parser.add_argument("input_files", nargs='+') |
|---|
| | | args = parser.parse_args() |
|---|
| | | |
|---|
| | | average_model(args.input_files, args.output_file) |
|---|
