File size: 556 Bytes
0b32ad6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import os
import sys
import yaml
import argparse
import torch
parser = argparse.ArgumentParser()
parser.add_argument('--ckpt_pth', required=True)
parser.add_argument('--save_dir', required=True)
paras = parser.parse_args()
ckpt = torch.load(paras.ckpt_pth, map_location='cpu')
os.makedirs(paras.save_dir, exist_ok=True)
with open(f'{paras.save_dir}/args.yaml', 'w') as handle:
args = ckpt['Args']
yaml.dump(vars(args), handle)
with open(f'{paras.save_dir}/config.yaml', 'w') as handle:
config = ckpt['Config']
yaml.dump(config, handle)
|