utils.py

import yaml
import os
import glob
import torch

def load_model(args):
    checkpoint_path = f'{os.path.abspath(os.path.dirname(__file__))}/checkpoints/{args.group}:{args.name}/*.ckpt'
    print("::: Loading model from", checkpoint_path)
    checkpoints = glob.glob(checkpoint_path)

    try:
        state_dict = torch.load(checkpoints[-1])
    except Exception as e:
        print("No checkpoints found: ", checkpoint_path)
        raise e

    return state_dict

def load_model_by_dir(name):
    checkpoint_path = f'{os.path.abspath(os.path.dirname(__file__))}/{name}/*.ckpt'
    print('::: Loading model from:', checkpoint_path)
    checkpoints = glob.glob(checkpoint_path)
    return torch.load(checkpoints[-1])

def load_model_by_name(name):
    checkpoint_path = f'{os.path.abspath(os.path.dirname(__file__))}/checkpoints/{name}/*.ckpt'
    checkpoints = glob.glob(checkpoint_path)
    return torch.load(checkpoints[-1])

def extend_config(cfg_path, child = None):
    with open(cfg_path, 'rt') as f:
        parent_cfg = yaml.load(f, Loader = yaml.FullLoader)

    if child is not None:
        parent_cfg.update(child)

    if '$extends$' in parent_cfg:
        path = parent_cfg['$extends$']
        del parent_cfg['$extends$']
        parent_cfg = extend_config(child = parent_cfg, cfg_path = path)

    return parent_cfg

def load_args(args):
    cfg = extend_config(cfg_path = f'{args.config_file}', child = None)

    for key, value in cfg.items():
        if key in args and args.__dict__[key] is not None:
            continue
        args.__dict__[key] = value

    return args, cfg