File size: 1,899 Bytes
961cf0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | import torch
from models import multimodalIntraInterModal
def _strip_module_prefix(state_dict):
keys = list(state_dict.keys())
if keys and keys[0].startswith("module."):
return {k.replace("module.","",1):v for k,v in state_dict.items()}
return state_dict
def _load_checkpoint(model_path, device):
try:
return torch.load(model_path, map_location=device)
except (RuntimeError, OSError) as exc:
msg = str(exc)
if "PytorchStreamReader failed reading zip archive" in msg:
raise RuntimeError(
f"Checkpoint file is not a readable PyTorch archive: {model_path}. "
"The file is likely truncated, corrupted, or was not exported correctly."
) from exc
raise RuntimeError(
f"Failed to load checkpoint from {model_path}: {exc}"
) from exc
def load_model(
device="cuda",
model_path:str=None,
num_classes=6,
cnn_model_name="densenet169",
attention_mecanism="gfcam",
vocab_size=91,
num_heads=8,
n=2,
text_model_name="one-hot-encoder",
unfreeze_weights="frozen_weights"
):
model = multimodalIntraInterModal.MultimodalModel(
num_classes=num_classes, device=device, cnn_model_name=cnn_model_name,
text_model_name=text_model_name, vocab_size=vocab_size, num_heads=num_heads,
attention_mecanism=attention_mecanism, n=n, unfreeze_weights=unfreeze_weights
)
ckpt = _load_checkpoint(model_path, device)
state_dict = ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt
model.load_state_dict(state_dict, strict=False)
model.to(device=device).eval()
return model
def find_last_conv(module):
last_conv = None
for m in module.modules():
if isinstance(m, torch.nn.Conv2d):
last_conv = m
return last_conv
|