def load_jit_model(arch_name, progress=True):
assert (
arch_name in PRETRAINED_MODELS
), f"Invalid arch {arch_name}, supported arch {PRETRAINED_MODELS.keys()}"
model_info = PRETRAINED_MODELS[arch_name]
model_path = model_info["model_path"]
if model_path.startswith("https://"):
model_path = hub_utils.download_file(model_path, progress=progress)
model = torch.jit.load(model_path, map_location="cpu")
model.model_info = model_info
return model
def _load_fbnet_state_dict(self, file_name, progress=True):
if file_name.startswith("https://"):
file_name = hub_utils.download_file(file_name, progress=progress)
state_dict = torch.load(file_name, map_location="cpu")["state_dict"]
ret = {}
for name, val in state_dict.items():
if name.startswith("module."):
name = name[len("backbone.module."):]
if name[0] == '.':
continue
ret[name] = val
return ret
def _load_fbnet_state_dict(file_name, progress=True):
if file_name.startswith("https://"):
file_name = hub_utils.download_file(file_name, progress=progress)
state_dict = torch.load(file_name, map_location="cpu")
if "model_ema" in state_dict and state_dict["model_ema"] is not None:
state_dict = state_dict["model_ema"]
else:
state_dict = state_dict["state_dict"]
ret = {}
for name, val in state_dict.items():
if name.startswith("module."):
name = name[len("module."):]
ret[name] = val
return ret