Function init_from_checkpoint

main.py:335–364

(cfg: DictConfig, new_model: nn.Module)

def init_from_checkpoint(cfg: DictConfig, new_model: nn.Module):
    print(f"Initializing model from checkpoint {cfg.checkpoint_run_name}")
    checkpoint_cfg = Path(cfg.checkpoint_cfg)
    assert checkpoint_cfg.exists(), f"Checkpoint config {checkpoint_cfg} does not exist"
    pretrained_cfg = om.load(checkpoint_cfg)

    pretrained_model = build_model(pretrained_cfg.model)
    n_params = sum(p.numel() for p in pretrained_model.parameters())

    checkpoint_filepath = Path(cfg.checkpoint_load_path) / f"{cfg.checkpoint_run_name}" / "latest-rank0.pt"
    assert checkpoint_filepath.exists(), f"Checkpoint {checkpoint_filepath} does not exist"
    state = torch.load(_ensure_valid_checkpoint(checkpoint_filepath), map_location="cpu")

    state_dict = state.get("state", {})
    model_state = state_dict.get("model", {})
    assert len(model_state) > 0, "Model state is empty, please check the checkpoint and checkpoint path"

    pretrained_model.load_state_dict(model_state)

    if isinstance(pretrained_cfg.model.model_config, DictConfig):
        model_config = OmegaConf.to_container(pretrained_cfg.model.model_config, resolve=True)
    pretrained_config = FlexBertConfig.from_pretrained(pretrained_cfg.model.pretrained_model_name, **model_config)

    init_mlm_model_from_pretrained(
        config=pretrained_config,
        pretrained_model=pretrained_model.model,
        new_model=new_model.model,
        mode=cfg.get("mode", "tile_weights_from_middle"),
    )
    print(f"Initalized model from checkpoint {cfg.checkpoint_run_name} with {n_params=:.4e} parameters")


def main(cfg: DictConfig, return_trainer: bool = False, do_train: bool = True) -> Optional[Trainer]: