def load_requirements(self, *args, **kwargs): if is_master(): requirements = self.config.get("zoo_requirements", []) if isinstance(requirements, str): requirements = [requirements] for item in requirements: download_pretrained_model(item, *args, **kwargs) synchronize()
def __init__(self, in_dim, weights_file, bias_file, model_data_dir, *args, **kwargs): super().__init__() model_data_dir = get_absolute_path(model_data_dir) if not os.path.isabs(weights_file): weights_file = os.path.join(model_data_dir, weights_file) if not os.path.isabs(bias_file): bias_file = os.path.join(model_data_dir, bias_file) if not PathManager.exists(bias_file) or not PathManager.exists( weights_file): download_path = download_pretrained_model("detectron.vmb_weights") weights_file = get_absolute_path( os.path.join(download_path, "fc7_w.pkl")) bias_file = get_absolute_path( os.path.join(download_path, "fc7_b.pkl")) with PathManager.open(weights_file, "rb") as w: weights = pickle.load(w) with PathManager.open(bias_file, "rb") as b: bias = pickle.load(b) out_dim = bias.shape[0] self.lc = nn.Linear(in_dim, out_dim) self.lc.weight.data.copy_(torch.from_numpy(weights)) self.lc.bias.data.copy_(torch.from_numpy(bias)) self.out_dim = out_dim
def _load_pretrained_model(model_name_or_path, *args, **kwargs): if PathManager.exists(model_name_or_path): download_path = model_name_or_path model_name = model_name_or_path else: download_path = download_pretrained_model(model_name_or_path, *args, **kwargs) model_name = model_name_or_path _hack_imports() ckpt_path = get_ckpt_path_from_folder(download_path) ckpt = get_ckpt_from_path(ckpt_path) # If configs are not present, will ckpt provide the config? config = get_config_from_folder_or_ckpt(download_path, ckpt) model_config = config.get("model_config", config) ckpt = ckpt.get("model", ckpt) # Also handle the case of model_name is path if PathManager.exists(model_name): # This shouldn't happen assert len(model_config.keys()) == 1, "Checkpoint contains more than one model?" # Take first key model_config = model_config[list(model_config.keys())[0]] else: model_config = model_config.get(model_name.split(os.path.sep)[-1].split(".")[0]) return {"config": model_config, "checkpoint": ckpt, "full_config": config}
def _load_pretrained_model(model_name_or_path, *args, **kwargs): if PathManager.exists(model_name_or_path): download_path = model_name_or_path model_name = model_name_or_path else: download_path = download_pretrained_model(model_name_or_path, *args, **kwargs) model_name = model_name_or_path configs = glob.glob(os.path.join(download_path, "*.yaml")) assert len(configs) <= 1, ( "Multiple yaml files with the pretrained model. " + "MMF doesn't know what to do.") ckpts = [] allowed_ckpt_types = [f"*{ext}" for ext in ALLOWED_CHECKPOINT_EXTS] for ckpt_type in allowed_ckpt_types: ckpts.extend(glob.glob(os.path.join(download_path, ckpt_type))) assert ( len(ckpts) == 1 ), "None or multiple checkpoints files. MMF doesn't know what to do." _hack_imports() with PathManager.open(ckpts[0], "rb") as f: ckpt = torch.load(f, map_location=lambda storage, loc: storage) # If configs are not present, will ckpt provide the config? if len(configs) == 0: assert "config" in ckpt, ( "No configs provided with pretrained model" " while checkpoint also doesn't have configuration.") config = ckpt["config"] else: config = load_yaml(configs[0]) model_config = config.get("model_config", config) ckpt = ckpt.get("model", ckpt) # Also handle the case of model_name is path if PathManager.exists(model_name): # This shouldn't happen assert len(model_config.keys() ) == 1, "Checkpoint contains more than one model?" # Take first key model_config = model_config[list(model_config.keys())[0]] else: model_config = model_config.get( model_name.split(os.path.sep)[-1].split(".")[0]) return {"config": model_config, "checkpoint": ckpt, "full_config": config}
def load_pretrained_model(model_name_or_path, *args, **kwargs): # If this is a file, then load this directly else download and load if PathManager.exists(model_name_or_path): download_path = model_name_or_path model_name = model_name_or_path else: download_path = download_pretrained_model(model_name_or_path, *args, **kwargs) model_name = model_name_or_path config_folder_path = os.path.dirname(download_path) configs = glob.glob(os.path.join(config_folder_path, "*.yaml")) assert len(configs) <= 1, ( "Multiple yaml files with the pretrained model. " + "MMF doesn't know what to do.") ckpts = [] allowed_ckpt_types = ("*.ckpt", "*.pth", "*.pt") for ckpt_type in allowed_ckpt_types: if download_path.endswith(ckpt_type.split("*")[1]): ckpts.extend(glob.glob(download_path)) else: ckpts.extend(glob.glob(os.path.join(download_path, ckpt_type))) assert ( len(ckpts) == 1 ), "None or multiple checkpoints files. MMF doesn't know what to do." _hack_imports() ckpt = torch.load(ckpts[0], map_location=lambda storage, loc: storage) # If configs are not present, will ckpt provide the config? if len(configs) == 0: assert "config" in ckpt, ( "No configs provided with pretrained model " " while checkpoint also doesn't have configuration.") config = ckpt["config"] else: config = load_yaml(configs[0]) model_config = config.get("model_config", config) ckpt = ckpt.get("model", ckpt) # Also handle the case of model_name is path model_config = model_config.get( model_name.split(os.path.sep)[-1].split(".")[0]) return {"config": model_config, "checkpoint": ckpt, "full_config": config}
def test_load_resume_zoo_parity_with_mmf(self): # not specifying checkpoint.resume, but specifying # checkpoint.resume_zoo. It should load the model file # underlying zoo resume_zoo = "unimodal_text.hateful_memes.bert" ckpt_filepath = download_pretrained_model(resume_zoo) ckpt_filepath = get_ckpt_path_from_folder(ckpt_filepath) ckpt = torch.load(ckpt_filepath, map_location="cpu") ckpt_config = {"resume_zoo": resume_zoo, "zoo_config_override": True} with mock_env_with_temp("mmf.utils.checkpoint.get_mmf_env") as _: mmf_trainer = self._get_mmf_trainer( ckpt_config=ckpt_config, model_config=unimodal_text_model_config, max_updates=0, ) mmf_trainer.on_init_start() mmf_ckpt = mmf_trainer.model.state_dict() mmf_ckpt.pop("base.encoder.embeddings.position_ids") self._assert_same_dict(ckpt, mmf_ckpt) with mock_env_with_temp( "mmf.trainers.lightning_trainer.get_mmf_env") as _: # lightning load from zoo, in this case, the zoo ckpt is in mmf format lightning = self._get_lightning_trainer( ckpt_config=ckpt_config, model_config=unimodal_text_model_config, max_steps=0, seed=4, ) lightning.trainer.fit(lightning.model, train_dataloaders=lightning.train_loader) lightning_ckpt = lightning.model.state_dict() lightning_ckpt.pop("base.encoder.embeddings.position_ids") self._assert_same_dict(ckpt, lightning_ckpt)
def load_requirements(self, config, *args, **kwargs): requirements = config.get("zoo_requirements", []) if isinstance(requirements, str): requirements = [requirements] for item in requirements: download_pretrained_model(item, *args, **kwargs)
def get_checkpoint_data(self) -> Dict[str, Any]: """This function gets checkpoint file path on disk from config.trainer.params.resume_from_checkpoint. However if it not specified, it gets checkpoint path from config.checkpoint. If config.resume is specified it gets the latest checkpoint from the config's save directory (alternatively it gets the best checkpoint if config.resume_best is True). If config.resume is not specified, then it gets config.resume_file or the checkpoint file from config.resume_zoo (in that order). Returns: Dict[str, Any]: a dict containing the following keys, `checkpoint_path` (str) local file path for the checkpoint; `ckpt` (Dict[str, Any]) `is_zoo` (Bool) whether or not the checkpoint is specified through a zoo identifier `config` (Dict[str, Any]]) the config that is stored together with this checkpoint """ # get ckpt file path from config.trainer.params.resume_from_checkpoint path = self.config.trainer.params.get("resume_from_checkpoint", None) if path is not None: is_zoo = self.is_zoo_path(path) ckpt_filepath = path if is_zoo: folder = download_pretrained_model(path) ckpt_filepath = get_ckpt_path_from_folder(folder) ckpt = get_ckpt_from_path(ckpt_filepath) config = get_config_from_folder_or_ckpt(folder, ckpt) else: ckpt = get_ckpt_from_path(ckpt_filepath) config = None return { "ckpt": ckpt, "checkpoint_path": ckpt_filepath, "is_zoo": is_zoo, "config": config, } is_zoo = False config = None ckpt = None # get ckpt file path from config.checkpoint ckpt_config = self.config.checkpoint suffix = "best.ckpt" if ckpt_config.resume_best else "current.ckpt" path = os.path.join(get_mmf_env(key="save_dir"), suffix) ckpt_filepath = None resume_from_specified_path = (ckpt_config.resume_file is not None or ckpt_config.resume_zoo is not None ) and (not ckpt_config.resume or not PathManager.exists(path)) if resume_from_specified_path: if ckpt_config.resume_file and PathManager.exists( ckpt_config.resume_file): ckpt_filepath = ckpt_config.resume_file elif ckpt_config.resume_zoo is not None: is_zoo = True folder = download_pretrained_model(ckpt_config.resume_zoo) ckpt_filepath = get_ckpt_path_from_folder(folder) ckpt = get_ckpt_from_path(ckpt_filepath) config = get_config_from_folder_or_ckpt(folder, ckpt) else: raise RuntimeError(f"{ckpt_config.resume_file} doesn't exist") if ckpt_config.resume and PathManager.exists(path): ckpt_filepath = path if ckpt_filepath is not None: ckpt = get_ckpt_from_path(ckpt_filepath) return { "ckpt": ckpt, "checkpoint_path": ckpt_filepath, "is_zoo": is_zoo, "config": config, }