feuermagier / beyond_deep_ensembles Goto Github PK

def load_model(model_idx, model, scaler, optimizer, out_path, config, log):
    ckpt = None
    start_epoch = 0

    # Load checkpoint and scaler if available
    if config.get("use_checkpoint", None):
        try:
            ckpt_paths = glob.glob(out_path + f"{config['model']}_chkpt_{model_idx}_*.pth")
            ckpt_paths.sort(key=os.path.getmtime)
            ckpt = torch.load(ckpt_paths[-1]) 
            model.load_state_dict(ckpt['model_state_dict'])
            start_epoch = ckpt["epoch"] + 1
            scaler.load_state_dict(ckpt["scaler_state_dict"])
            log.info(f"Loaded checkpoint for model {model_idx} at epoch {start_epoch}")
        except:
            log.info(f"Failed to load checkpoint for model {model_idx}")

    optimizer.init_grad_scaler(scaler)

    # Load optimizer state if available
    # Base optimizer state is loaded separately if available
    if ckpt is not None: 
        try:
            optimizer.load_state_dict(ckpt["optimizer_state_dict"])
            if ckpt.get("base_optimizer") is not None:
                optimizer.get_base_optimizer().load_state_dict(ckpt["base_optimizer"])
            log.info(f"Loaded base optimizer state for model {model_idx}")
        except:
            log.info(f"Failed to load optimizer state for model {model_idx}")

    # Load scheduler state if available
    if config["lr_schedule"]:
        scheduler = wilson_scheduler(optimizer.get_base_optimizer(), config["epochs"], config["lr"], None)
        if ckpt is not None:
            scheduler.load_state_dict(ckpt["scheduler_state_dict"])
            log.info(f"Loaded scheduler state for model {model_idx}")
    else:
        scheduler = None
    
    return start_epoch, model, optimizer, scaler, scheduler

def save_model(model, optimizer, scheduler, scaler, out_path, config, model_idx, epoch):
    state_dict = {
                    'epoch': epoch,
                    'model_idx': model_idx,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scaler_state_dict': scaler.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict() if scheduler is not None else 'None'
                }
    if hasattr(optimizer, "get_base_optimizer"):
        state_dict['base_optimizer'] = optimizer.get_base_optimizer().state_dict()
    torch.save(state_dict, out_path + f"{config['model']}_chkpt_{model_idx}_{epoch}.pth")

class MAPOptimizer(BayesianOptimizer):
    '''
        Maximum A Posteriori

        This simply optimizes a point estimate of the parameters with the given base_optimizer.
    '''

    def __init__(self, params, base_optimizer):
        super().__init__(params, {})
        # self.state["__base_optimizer"] = base_optimizer
        self.__base_optimizer = base_optimizer