shithub: opus

ref: b9912f7eb3455926987434db5e3dc97e1a24c1e9
dir: /dnn/torch/osce/extract_setup.py/

View raw version 
import torch
import yaml
import argparse


parser = argparse.ArgumentParser()
parser.add_argument('checkpoint', type=str, help='model checkpoint')
parser.add_argument('setup', type=str, help='setup filename')

if __name__ == "__main__":
    args = parser.parse_args()

    ckpt = torch.load(args.checkpoint, map_location='cpu')

    setup = ckpt['setup']

    with open(args.setup, "w") as f:
        yaml.dump(setup, f)