def _restore_model_weights(self, model):
"""
If using a weights file to initialize the model, we load the weights
and initialize the model. Since the weights file specified
by user might not be VISSL trained weights, we expose several config
options like APPEND_PREFIX, etc to allow successful loading of the weights.
See MODEL.WEIGHTS_INIT description in vissl/config/defaults.yaml for details.
"""
params_from_file = self.config["MODEL"]["WEIGHTS_INIT"]
init_weights_path = params_from_file["PARAMS_FILE"]
assert init_weights_path, "Shouldn't call this when init_weight_path is empty"
logging.info(f"Initializing model from: {init_weights_path}")
if PathManager.exists(init_weights_path):
weights = load_and_broadcast_checkpoint(init_weights_path,
device=torch.device("cpu"))
skip_layers = params_from_file.get("SKIP_LAYERS", [])
replace_prefix = params_from_file.get("REMOVE_PREFIX", None)
append_prefix = params_from_file.get("APPEND_PREFIX", None)
state_dict_key_name = params_from_file.get("STATE_DICT_KEY_NAME",
None)
# we initialize the weights from this checkpoint. However, we
# don't care about the other metadata like iteration number etc.
# So the method only reads the state_dict
init_model_from_weights(
self.config,
model,
weights,
state_dict_key_name=state_dict_key_name,
skip_layers=skip_layers,
replace_prefix=replace_prefix,
append_prefix=append_prefix,
)
return model
def build_retrieval_model(cfg):
"""
Builds the model on 1-gpu and initializes from the weight.
"""
logging.info("Building model....")
model = build_model(cfg.MODEL, cfg.OPTIMIZER)
if PathManager.exists(cfg.MODEL.WEIGHTS_INIT.PARAMS_FILE):
init_weights_path = cfg.MODEL.WEIGHTS_INIT.PARAMS_FILE
logging.info(f"Initializing model from: {init_weights_path}")
weights = torch.load(init_weights_path, map_location=torch.device("cuda"))
skip_layers = cfg.MODEL.WEIGHTS_INIT.get("SKIP_LAYERS", [])
replace_prefix = cfg.MODEL.WEIGHTS_INIT.get("REMOVE_PREFIX", None)
append_prefix = cfg.MODEL.WEIGHTS_INIT.get("APPEND_PREFIX", None)
state_dict_key_name = cfg.MODEL.WEIGHTS_INIT.get("STATE_DICT_KEY_NAME", None)
init_model_from_weights(
cfg,
model,
weights,
state_dict_key_name=state_dict_key_name,
skip_layers=skip_layers,
replace_prefix=replace_prefix,
append_prefix=append_prefix,
)
else:
# We only throw the warning if not weights file is provided. We want to
# benchmark the random initialization model too and hence support that.
logging.warning("Model is randomly initialized....")
logging.info(f"Model is:\n {model}")
return model