DeepECG-SSL fine-tuned on CLSA for incident CVD (5 fold models, Model A+ parameters, inference)
7916811 verified | """Inference for DeepECG-SSL-CLSA-CVD (Model A / Model A+). | |
| Requires: torch, safetensors, numpy, fairseq-signals (https://github.com/Jwoo5/fairseq-signals) and the public | |
| DeepECG-SSL v2 (amplitude-preserved) pretrained encoder checkpoint, which defines the architecture. | |
| ECG input: (N, 12, 2500) float32 in millivolts, 250 Hz, 10 s, standard lead order, no per-lead normalisation. | |
| from inference import DeepECGSSLCVD | |
| m = DeepECGSSLCVD("weights/", "path/to/deepecg-ssl-v2-wcr-amp-preserved-encoder.pt", device="cuda") | |
| p_ecg = m.ecg_probability(ecg_mv) # Model A (calibrated ECG-only probability) | |
| risk = m.model_a_plus(ecg_mv, age, female, hypertension) # Model A+ (no blood test) | |
| """ | |
| import json, os, numpy as np, torch, torch.nn as nn | |
| from safetensors.torch import load_file | |
| from fairseq_signals.utils import checkpoint_utils | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| lg = lambda p: np.log(np.clip(p, 1e-6, 1 - 1e-6) / (1 - np.clip(p, 1e-6, 1 - 1e-6))) | |
| class Net(nn.Module): | |
| def __init__(self, base_ckpt): | |
| super().__init__() | |
| models, _, _ = checkpoint_utils.load_model_and_task(base_ckpt) | |
| self.enc = models[0] if isinstance(models, list) else models | |
| for n in ["quantizer", "project_q", "final_proj", "target_glu"]: | |
| if hasattr(self.enc, n): setattr(self.enc, n, None) | |
| self.head = nn.Sequential(nn.Dropout(0.3), nn.Linear(768, 1)) | |
| def forward(self, x): | |
| return self.head(self.enc(source=x, mask=False, features_only=True)["x"].mean(1)).squeeze(1) | |
| class DeepECGSSLCVD: | |
| def __init__(self, weights_dir, base_ckpt, device="cpu"): | |
| self.p = json.load(open(os.path.join(HERE, "model_params.json"))); self.device = device; self.models = [] | |
| for k in range(5): | |
| m = Net(base_ckpt); m.load_state_dict(load_file(os.path.join(weights_dir, f"deepecg_ssl_clsa_cvd_fold{k}.safetensors")), strict=False) | |
| self.models.append(m.to(device).eval()) | |
| def ecg_probability(self, ecg_mv, batch=64): | |
| x = torch.as_tensor(np.asarray(ecg_mv, dtype=np.float32)); out = [] | |
| for k, m in enumerate(self.models): | |
| pk = [] | |
| for i in range(0, len(x), batch): | |
| pk.append(torch.sigmoid(m(x[i:i + batch].to(self.device))).float().cpu().numpy()) | |
| a, b = self.p["per_fold_platt"][str(k)]["a"], self.p["per_fold_platt"][str(k)]["b"] | |
| out.append(1 / (1 + np.exp(-(a * lg(np.concatenate(pk)) + b)))) | |
| return np.mean(out, 0) | |
| def model_a_plus(self, ecg_mv, age, female, hypertension): | |
| c = self.p["model_A_plus"]; e = lg(self.ecg_probability(ecg_mv)) | |
| z = c["intercept"] + c["b_ecg"] * e + c["b_age"] * np.asarray(age, float) + c["b_female"] * np.asarray(female, float) + c["b_htn"] * np.asarray(hypertension, float) | |
| return 1 / (1 + np.exp(-z)) | |