| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class InferenceGraphSAGE(nn.Module): |
| def __init__(self, in_dim: int = 768, hidden_dim: int = 256, num_layers: int = 2, dropout: float = 0.1): |
| super().__init__() |
| self.dropout = float(dropout) |
|
|
| self.linears = nn.ModuleList() |
| cur_in = in_dim |
| for _ in range(num_layers): |
| self.linears.append(nn.Linear(cur_in * 2, hidden_dim)) |
| cur_in = hidden_dim |
|
|
| self.class_wheeze = nn.Linear(cur_in, 1) |
| self.class_crackle = nn.Linear(cur_in, 1) |
| self.act = nn.ReLU() |
|
|
| def forward(self, x: torch.Tensor, adj: torch.Tensor): |
| for lin in self.linears: |
| deg = adj.sum(dim=-1, keepdim=True).clamp(min=1.0) |
| neighbor_mean = torch.matmul(adj, x) / deg |
| h = torch.cat([x, neighbor_mean], dim=-1) |
| h = self.act(lin(h)) |
| h = F.dropout(h, p=self.dropout, training=self.training) |
| x = h |
|
|
| wheeze = torch.sigmoid(self.class_wheeze(x).squeeze(-1)) |
| crackle = torch.sigmoid(self.class_crackle(x).squeeze(-1)) |
| return wheeze, crackle |
|
|
|
|
| def load_checkpoint_flexible(model: nn.Module, path: str, map_location: Optional[str] = "cpu") -> nn.Module: |
| ck = torch.load(path, map_location=map_location) |
| if isinstance(ck, dict) and ("state_dict" in ck or "model_state_dict" in ck): |
| state = ck.get("state_dict", ck.get("model_state_dict")) |
| else: |
| state = ck |
|
|
| try: |
| model.load_state_dict(state, strict=False) |
| return model |
| except Exception: |
| ms = model.state_dict() |
| incoming = {} |
| for k, v in state.items(): |
| k2 = k[7:] if k.startswith("module.") else k |
| incoming[k2] = v |
|
|
| resolved = {} |
| for k, v in incoming.items(): |
| if k in ms and ms[k].shape == v.shape: |
| resolved[k] = v |
|
|
| for k_in, v in incoming.items(): |
| if k_in in resolved: |
| continue |
| for k_model, v_model in ms.items(): |
| if k_model in resolved: |
| continue |
| if v_model.shape == v.shape and k_model.endswith(k_in.split(".")[-1]): |
| resolved[k_model] = v |
| break |
|
|
| ms.update(resolved) |
| model.load_state_dict(ms) |
| return model |
|
|
