def main(args):
    config = yaml_util.load_yaml_file(args.config)
    if (args.json is not None):
        main_util.overwrite_config(config, args.json)
    (distributed, device_ids) = main_util.init_distributed_mode(args.world_size, args.dist_url)
    device = torch.device((args.device if torch.cuda.is_available() else 'cpu'))
    teacher_model = get_model(config['teacher_model'], device)
    module_util.freeze_module_params(teacher_model)
    student_model_config = config['student_model']
    student_model = get_model(student_model_config, device)
    freeze_modules(student_model, student_model_config)
    ckpt_file_path = config['student_model']['ckpt']
    train_config = config['train']
    optim_config = train_config['optimizer']
    optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params'])
    scheduler_config = train_config['scheduler']
    lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params'])
    if file_util.check_if_exists(ckpt_file_path):
        (best_val_map, _, _) = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler)
        save_ckpt(student_model, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path)
def main(args):
    config = yaml_util.load_yaml_file(args.config)
    if (args.json is not None):
        main_util.overwrite_config(config, args.json)
    (distributed, device_ids) = main_util.init_distributed_mode(args.world_size, args.dist_url)
    device = torch.device((args.device if torch.cuda.is_available() else 'cpu'))
    teacher_model = get_model(config['teacher_model'], device)
    module_util.freeze_module_params(teacher_model)
    student_model_config = config['student_model']
    student_model = get_model(student_model_config, device)
    freeze_modules(student_model, student_model_config)
    ckpt_file_path = config['student_model']['ckpt']
    train_config = config['train']
    optim_config = train_config['optimizer']
    optimizer = func_util.get_optimizer(student_model, optim_config['type'], optim_config['params'])
    scheduler_config = train_config['scheduler']
    lr_scheduler = func_util.get_scheduler(optimizer, scheduler_config['type'], scheduler_config['params'])
    if file_util.check_if_exists(ckpt_file_path):
        (best_val_map, _, _) = load_ckpt(ckpt_file_path, optimizer=optimizer, lr_scheduler=lr_scheduler)
        save_ckpt(student_model, optimizer, lr_scheduler, best_val_map, config, args, ckpt_file_path)