| import sys | |
| import yaml | |
| import argparse | |
| import torch | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--ckpt_pth', required=True) | |
| parser.add_argument('--settings_dir', required=True) | |
| parser.add_argument('--new_ckpt_pth', required=True) | |
| paras = parser.parse_args() | |
| ckpt = torch.load(paras.ckpt_pth, map_location='cpu') | |
| with open(f'{paras.settings_dir}/args.yaml', 'r') as handle: | |
| args = yaml.load(handle, Loader=yaml.FullLoader) | |
| ckpt['Args'] = argparse.Namespace(**args) | |
| with open(f'{paras.settings_dir}/config.yaml', 'r') as handle: | |
| config = yaml.load(handle, Loader=yaml.FullLoader) | |
| ckpt['Config'] = config | |
| torch.save(ckpt, paras.new_ckpt_pth) | |