File size: 672 Bytes
0b32ad6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
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)
|