| 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 |
|
|
