Repository navigation
Conversation
Test CoverageTotal: 73.70% (20055 / 27210 lines)
|
…lassification --- Summary - PTBXLDataset — wraps the open-access PTB-XL 12-lead ECG dataset (21,799 records, 18,869 patients). Reads ptbxl_database.csv + scp_statements.csv, aggregates SCP codes into six label sets (all, diagnostic, subdiagnostic, superdiagnostic, form, rhythm) and writes a derived metadata CSV so downstream tasks can load labels without re-parsing the raw files. - PTBXLMultilabelClassification — task that converts each labelled ECG into a (12, 10 × fs) float32 signal with multi-hot labels. Supports 100 Hz and 500 Hz, likelihood filtering, per-record z-score normalisation, and optional 0.5 Hz high-pass baseline removal. Ships a fallback WFDB reader so wfdb is not a hard dependency. - ProtoECGNet — re-implementation of Sethi et al., MLHC 2025 (https://arxiv.org/abs/2504.08713). A 1D or 2D ResNet extracts features; each prototype branch computes cosine or L2 similarities between feature windows and learned prototype vectors assigned to labels; a final linear layer produces multi-label logits. Prototypes can be projected onto real ECG segments after training, making every prediction explainable by example cases. Supports multi-branch setups (the paper's rhythm / morphology / global split), learned fusion, contrastive and diversity regularisation, and the paper's four-stage training protocol. - Docs — RST stubs wired into the existing API reference pages. - Example script — end-to-end pipeline with ablation study (contrastive loss on/off × cosine/L2 × prototype count); runs on synthetic data with --synthetic (no download needed). Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
db1bc73 to
621377f
Compare
Logiquo
added
the
status: automatic
|
/review |
1 similar comment
|
/review |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Previous review
AI review: findings
The PR adds a PTB-XL dataset (builds a metadata CSV from ptbxl_database.csv and scp_statements.csv), a multi-label ECG task with a WFDB fallback reader, and ProtoECGNet (1D/2D ResNet prototype branches, contrastive/diversity/separation losses, staged training and prototype projection), plus docs, an example script and 45 synthetic-data unit tests. Nothing in the repo already provides a PTB-XL loader, a prototype model, or a 1D ResNet; CardiologyDataset targets the PhysioNet 2020 challenge with different labels, and the 2D backbone reuses torchvision ResNet-18, so no redundancy was found. The change only adds classes and exports and modifies no base class or existing behavior, so it is not core. Blocking (medium): in ProtoECGNet.set_stage (pyhealth/models/protoecgnet.py:707), frozen backbone parts are not put in eval mode, and Trainer.train calls model.train() every epoch (trainer.py:246), so BatchNorm running stats keep updating during the classifier and fusion stages after prototypes were projected. Blocking (medium): with fusion="learned" (the --paper-groups setting), predictions use branch logits in every stage except "fusion" (protoecgnet.py:783), and prototype_contributions (protoecgnet.py:898) reports branch contributions that do not sum to the output logit, so the explanations do not match the predictions. Non-blocking: the 500 Hz path and remove_baseline have only shape-level or no tests, the metadata CSV is written into the raw data root and never regenerated when label definitions change, and project_prototypes describes projected global prototypes as ECG segments even though the default global mode stores time-averaged features. Overall risk is moderate: parsing, label aggregation and loss math are covered by hand-checked tests, but the two blocking issues affect the staged protocol and the explanation API.
Reviewed commit 621377f
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a PTB-XL dataset, a multi-label ECG task and a ProtoECGNet re-implementation, with docs, an example script and a thorough synthetic-data test suite. Test coverage is above the bar: 97–98% on the three new modules, with the repository total at 68.45% against the 60% bar. You're close, but two problems in the staged protocol and the explanations need fixing first.
To get this merged:
- BatchNorm drifts in the frozen stages.
set_stageonly togglesrequires_grad.Trainer.traincallsmodel.train()every epoch, so the backbone and add-on BatchNorm running statistics keep updating during theclassifierandfusionstages. That shifts the features after the prototypes have been projected onto real segments, so the projected prototypes no longer match the features they were taken from. Please keep the frozen modules in eval mode for those stages. One way is to overridetrain()so that frozen backbone and add-on modules stay ineval(). Please also add a test that the BN running stats don't change during theclassifierandfusionstages. - Explanations don't match predictions with
fusion="learned". In every stage exceptfusion, the output logits come from the per-branch classifiers. In thefusionstage they come fromfusion_classifier.prototype_contributions()always reports similarity times the branch classifier weight, so with learned fusion (the--paper-groupssetting) the contributions don't sum to the model's logits. Please makeprototype_contributions()use the fusion weights when the stage isfusionor fusion is learned, and add a test that the contributions sum to the logit. If you'd rather keep it as is, document the limitation clearly.
Optional suggestions:
- The
project_prototypesdocstring calls the result "real ECG segments". For global mode (proto_time_len=None), the stored vector is a time-averaged feature, so please say so. - Add a test for the 500 Hz path and for
remove_baselineoutput values. - The metadata CSV is written into the raw data root and never regenerated if label definitions change. Consider an option to force regeneration, or write it to the cache directory.
Reviewed commit 621377f
Blocking fixes:
- Override train() to keep frozen backbone/add-on BatchNorm layers in
eval mode during classifier and fusion stages; frozen module set is
computed once in set_stage so train() has no per-call allocation
- Fix prototype_contributions() to use fusion_classifier weights when
stage="fusion" so contributions sum to the logit (branch classifier
weights were used unconditionally before)
Optional:
- Add refresh_metadata parameter to PTBXLDataset for forced CSV rebuild
- Add MPS device support in the example script for Apple Silicon
- Clarify project_prototypes docstring: global mode stores a
time-averaged feature vector, not a raw ECG segment
Tests:
- Verify BN running stats are unchanged during classifier and fusion
stages after a train()+forward pass
- Verify prototype_contributions sums to logit minus fusion bias
- Add 500 Hz path test and remove_baseline DC-suppression value test
- Add refresh_metadata round-trip test
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
There was a problem hiding this comment.
Previous review
AI review: findings
The PR adds a PTB-XL dataset (pyhealth/datasets/ptbxl.py), a multi-label ECG task with a fallback WFDB reader (pyhealth/tasks/ptbxl_multilabel_classification.py), and ProtoECGNet (pyhealth/models/protoecgnet.py), a prototype model with staged training, prototype projection and per-prototype contributions. It builds on BaseDataset, BaseTask, BaseModel and Trainer without modifying them, and only adds exports and docs. No existing PyHealth model or interpretability method implements prototypes or a 1D/2D ResNet backbone, so nothing is redundant; CardiologyDataset already covers PTB-XL recordings via the Challenge 2020 bundle, but with a different label scheme and is not a duplicate. The code is mechanically consistent with the base classes: the cache keys on source fingerprints, set_task keys on task attributes, and the trainer's monitor and stage loop behave as the example expects. Tests use synthetic data and cover the loss terms, shapes, stage freezing, projection and fusion, but the PR states 45 tests while the files contain 51, and the example is never run in CI. Blocking (medium): the "every prediction can be explained by example cases" claim is not supported by the model, because project_prototypes returns the training-sample index and window only to the caller and nothing stores them, so they are lost after save/load, and prototype_contributions returns weights only (protoecgnet.py:894, :897). Blocking (medium): the paper command --label-set all --paper-groups has no reported run, and sklearn's macro ROC AUC raises if any of the 71 labels has no positives in the validation fold (example:222), so this advertised path may crash. Non-blocking (low): fusion="average" is only a scatter (protoecgnet.py:567), the fusion contribution docstring omits the bias (:908), torchvision's unused fc and avgpool stay in ResNet2dBackbone (:208-211), the example leaves a multi-GB temp cache per run (example:285), baked-in metadata paths can break with a relative root (ptbxl.py:179), and L1 is off by default in the example (example:260). Overall risk is moderate: the code is sound for the 5-label superdiagnostic setting that was run, but the paper-setting path and the explanation claim are unverified.
Reviewed commit 9e28dcb
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a PTB-XL dataset, a multi-label ECG task and a ProtoECGNet re-implementation, with thorough tests. Coverage is 69.59% for the whole repository, against a 60% bar. The new files are at 98.3% (ptbxl.py), 97.2% (protoecgnet.py) and 95.6% (ptbxl_multilabel_classification.py), against an 80% bar. The work is close, but two issues need attention first.
To get this merged:
- Paper-setting command (
--label-set all --paper-groups --paper-prototypes).- It is advertised, but no run is reported for it.
- Macro ROC AUC from sklearn raises an error if any of the 71 labels has no positive samples in the validation or test fold. Several rare labels may hit this.
- Please run this path, or guard against it. For example, compute macro AUROC only over labels that have positives, or warn and skip those labels.
- If you can't verify the path, remove the claim that it is supported.
- "Every prediction can be explained by example cases."
project_prototypes()returns the source sample index and window only to the caller. Nothing is stored on the model, so the information is lost after save/load.prototype_contributions()returns weights only.- Please either persist the projection info, for example in a buffer or attribute included in
state_dict, or soften the claim in the PR description and docstrings. A softened claim could say that projection returns the source segments to the caller.
Optional suggestions:
- The description says 45 tests, but the files contain about 51. Please update the count.
fusion="average"is only a scatter of the branch logits. A docstring note would make that clearer.- The fusion contribution docstring omits the bias term.
ResNet2dBackbonekeeps torchvision's unusedfcandavgpool.- The example leaves a temporary cache directory behind on every run.
- The metadata CSV stores paths built from
root, so a relativerootcan break later. Consider usingos.path.abspath. - The example leaves L1 off by default (
--lam-l1 0).
Reviewed commit 9e28dcb
- prototype_contributions docstring: clarify that in fusion mode
contributions sum to logit minus fusion_classifier.bias (fusion
classifier has a bias; branch classifiers do not)
- ResNet2dBackbone: replace unused avgpool and fc with nn.Identity()
to drop ~512K dead parameters
- Example: wrap dataset/training in try/finally so the scratch
directory (potentially multi-GB litdata cache) is always cleaned up;
move parser.error() calls before mkdtemp so they exit cleanly
- Example: expand --lam-l1 help text to explain when to enable it
- ptbxl.py: use os.path.abspath(root) in prepare_metadata so
record paths in the metadata CSV survive os.chdir
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
- prototype_contributions docstring: clarify that in fusion mode
contributions sum to logit minus fusion_classifier.bias (fusion
classifier has a bias; branch classifiers do not)
- ResNet2dBackbone: replace unused avgpool and fc with nn.Identity()
to drop ~512K dead parameters
- Example: wrap dataset/training in try/finally so the scratch
directory (potentially multi-GB litdata cache) is always cleaned up;
move parser.error() calls before mkdtemp so they exit cleanly
- Example: add --cache-dir option so repeated real-data runs reuse
the processed cache instead of re-preprocessing several GB each time
- Example: expand --lam-l1 help text to explain when to enable it
- ptbxl.py: use os.path.abspath(root) in prepare_metadata so
record paths in the metadata CSV survive os.chdir
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
There was a problem hiding this comment.
Previous review
AI review: findings
Adds a PTB-XL dataset loader (PTBXLDataset), a 12-lead multilabel task with a fallback WFDB reader, a ProtoECGNet model (1D/2D ResNet prototype branches with cosine or L2 similarity, staged training, projection onto real training segments), RST docs, a pipeline/ablation example, and 51 unit tests on synthetic tensors and folders. All changes are additive: three init export lines and new files, with no base class, shared helper, default, or existing API modified. Redundancy: PyHealth has no other prototype model and no other WFDB reader, so the model and task are not redundant; PTBXLDataset only partly overlaps CardiologyDataset's Challenge-2020 ptb-xl subset (different labels and folds), and the PR does not acknowledge that (low). Blocking (medium): pyhealth/datasets/ptbxl.py:187 writes ptbxl-metadata-pyhealth.csv into the raw input root on first construction, which fails on read-only inputs such as Kaggle's /kaggle/input and is not disclosed in the PR. Non-blocking (low): lam_l1 defaults to 0, so the default run skips the L1 step the docstring describes as from the paper; the contrastive loss normalisation (protoecgnet.py:479) is not listed among the differences from the reference; ProtoECGNet silently ignores BaseModel.set_pos_weight because it has its own class_weights; the example selects MPS when available, contradicting the description's CPU-only claim; and project_prototypes silently mis-indexes projection_info if given a shuffled loader. I did not run the tests or the example, and real-data AUROC numbers cannot be checked from the repository. Overall risk is low to moderate: the model is isolated, and the main concern is the write into the input directory.
Reviewed commit b9a9d3a
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a PTB-XL dataset, a multi-label ECG task with a fallback WFDB reader, and a ProtoECGNet implementation with staged training, prototype projection, docs, an example script and 51 tests. Coverage is good. The whole repo is at 70.7%, against a bar of 60%. The new files are at 98.3% (ptbxl.py), 97.2% (protoecgnet.py) and 95.6% (ptbxl_multilabel_classification.py), against a bar of 80%. I found one blocking issue.
To get this merged:
- PTBXLDataset.init and prepare_metadata (pyhealth/datasets/ptbxl.py, around line 187) write ptbxl-metadata-pyhealth.csv into the raw data
root. This fails on read-only inputs such as a Kaggle /kaggle/input mount or a shared read-only data directory. The PR description doesn't mention it. Please either:- write the derived CSV to
cache_diror another writable location, and fall back to it whenrootisn't writable, or - catch the OSError and give a clear message.
Whichever you choose, document it in the docstring. Updatesplit_by_foldin the example script to find the file in its new location.
- write the derived CSV to
- Add a test for the fallback or error path you choose.
Optional suggestions:
- project_prototypes records
sampleindices by loader order, so a shuffled loader would silently give wrong indices. Consider raising an error or documenting this more prominently. - lam_l1 defaults to 0, so the default run skips the L1 step. Please state this next to the "Differences from the paper" list. Also add the contrastive-loss normalisation (the division by count**0.5) to that list.
- ProtoECGNet uses its own class_weights and ignores BaseModel.set_pos_weight. A short docstring note would help.
- The example script picks MPS when it is available, but the description says CPU only. Please align the two.
- PTBXLDataset partly overlaps the PTB-XL subset in CardiologyDataset. A short note on how they differ would help users choose between them.
Reviewed commit b9a9d3a
When cache_dir is provided, ptbxl-metadata-pyhealth.csv is now written there rather than into root, so PTBXLDataset works on Kaggle /kaggle/input mounts and other read-only data directories. A patched YAML config with a relative path is generated in cache_dir so BaseDataset can still locate the file. self.metadata_path exposes the resolved path to callers. Also adds: shuffled-loader guard in project_prototypes, CardiologyDataset overlap note in PTBXLDataset docstring, set_pos_weight note in ProtoECGNet docstring, and one new test for each guard. Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
1 similar comment
|
/review |
There was a problem hiding this comment.
Previous review
AI review: findings
The PR adds a PTB-XL dataset (PTBXLDataset plus a derived metadata CSV), a multi-label ECG task with a fallback WFDB reader, and ProtoECGNet (1D/2D ResNet prototype branches, staged training, projection, and losses), with exports, RST docs, a synthetic-data example, and synthetic-data tests. It is self-contained: it subclasses BaseDataset, BaseTask and BaseModel without changing them, and the package init edits only add exports. I found no redundancy: PyHealth has no PTB-XL loader, no prototype model, no WFDB reader and no 1D residual backbone, and CardiologyDataset reads the Challenge 2020 .mat files. Blocking (medium): pyhealth/datasets/ptbxl.py:230 writes absolute record paths into the metadata CSV, and line 132 reuses an existing CSV without checking the root, so moving the PTB-XL folder or pointing cache_dir at another root makes set_task fail with FileNotFoundError. Blocking (medium): examples/ptbxl_multilabel_classification_protoecgnet.py:313 selects Apple MPS when available, which contradicts the PR's "CPU only" claim for the MacBook results. Non-blocking (low): ProtoECGNet silently ignores the inherited set_pos_weight (documented in the class Note), and lam_l1 defaults to 0, so the default pipeline omits the L1 step that the PR's differences list describes as part of the method. Tests cover the dataset, task, WFDB reader, losses, stage freezing, projection and fusion on synthetic data, which matches repository practice; the example and real-data path are not exercised in CI. Core behaviour I checked is consistent (Trainer.inference returns [y_true, y_prob, loss], projection index mapping is deterministic, and the cache key includes task parameters). Overall risk is low to moderate, and the blocking items are small fixes.
Reviewed commit 8916bba
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a PTB-XL dataset, a multi-label ECG task with a WFDB fallback reader, and a ProtoECGNet re-implementation with staged training and prototype projection. The tests are thorough, and coverage for the three new files is well above the bar: 97% for ptbxl.py, 97% for protoecgnet.py and 96% for the task file. Repository-wide coverage is 70.7%, above the 60% bar. It's close, but there are two things to fix first.
To get this merged:
- Stale or moved-root metadata.
prepare_metadatawrites absolute record paths intoptbxl-metadata-pyhealth.csv.__init__then reuses an existing CSV without checking that it matchesroot. If the PTB-XL folder is moved, or acache_diris reused with a different root,set_taskfails withFileNotFoundError. Please do one of the following:- Store paths relative to
rootand join them withrootwhen the task reads them. - Or record the root in the CSV, or beside it, and regenerate the file when it doesn't match.
Please also add a test for this case.
- Store paths relative to
- CPU-only claim versus the code. The description says the MacBook runs were CPU only and that the example script does not use the Apple GPU. But
main()selectsmpswhenever it is available. Please either make the script match the description, for example by adding a--deviceflag that defaults to CPU or CUDA, or correct the description and say which device each result came from.
Optional suggestions:
lam_l1defaults to 0, so the default pipeline skips the L1 step that the "Differences from the paper" section lists as part of the method. Consider mentioning this in the description, or enabling it in the paper-style settings.- Add a small test or doc note for
set_pos_weightbeing ignored by this model.
Reviewed commit 8916bba
…r dataset API Resolves the add/add conflict between our branch and the PTB-XL dataset implementation merged in sunlabuiuc#1187 (AxelNoun). Our dataset implementation is replaced by the one in master, which is strictly better in every way: - Atomic writes under FileLock (safe for DDP / parallel workers) - Content-hash cache keys (ptbxl_database.csv SHA1 in dataset_name, so the cache auto-invalidates when the source changes; no manual refresh_metadata= flag needed) - HIPAA age-censor handling (age=300 sentinel → age_is_censored flag) - Clean data_root / cache separation via load_table() override; the metadata CSV never touches the user's data root - sampling_rate as a dataset constructor parameter; a single signal_file attribute per event (no record_path_lr / record_path_hr pair) What this PR still adds on top of sunlabuiuc#1187: • PTBXLMultilabelClassification task: all six SCP-ECG label sets (all 71 statements / diagnostic / subdiagnostic / superdiagnostic / form / rhythm) with per-lead z-scoring and optional 0.5 Hz baseline removal. Re-wired to master's PTBXLDataset (event_type="records", event.signal_file, event.scp_codes, event.sampling_rate). scp_statements.csv aggregation is now done on-the-fly in the task using a module-level functools.cache loader; no pre-computed label columns needed. • ProtoECGNet model: case-based interpretable multi-label ECG classifier with staged training, prototype projection, learned fusion, and BatchNorm freeze. Unchanged from the previous review round (all 38 unit tests pass). • Full pipeline example: synthetic and real-data modes; uses split_by_strat_fold from master's splitter module. • 8 new tests in TestPTBXLMultilabelClassification covering signal IO, label aggregation with likelihood filtering, and the full set_task path against the master fixture. References PR sunlabuiuc#1187 (merged) and PR sunlabuiuc#1323 (this PR). Merge remote-tracking branch 'origin/master' into protoecgnet-ptbxl Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
There was a problem hiding this comment.
Previous review
AI review: findings
Adds PTBXLMultilabelClassification (six SCP label sets, a fallback WFDB reader and signal preprocessing), ProtoECGNet (a 1D or 2D ResNet prototype model with staged training, prototype projection and explanation helpers), an example script, RST docs and 51 synthetic-data tests. Base classes and existing APIs are unchanged; the new exports are additive. Blocking problems: (1) examples/ptbxl_multilabel_classification_protoecgnet.py:336 calls split_by_strat_fold(samples, folds=[range(1, 9), [9], [10]]), but splitter.py:442-447 requires one fold integer per sample and raises ValueError, so the example fails on every run, including --synthetic. (2) pyhealth/models/protoecgnet.py:766 (train override, with the frozen set built in set_stage at line 756) puts the whole PrototypeBranch into eval mode during the feature_extractor stage, because its own prototype_vectors are frozen, so the trainable backbone's BatchNorm never updates its running stats; no test covers this. (3) Redundant helpers: the task's read_wfdb_signal (ptbxl_multilabel_classification.py:25) duplicates pyhealth/datasets/ptbxl.py load_ptbxl_record, and its superdiagnostic aggregation (_aggregate_labels, line 96) duplicates PTBXLSuperclassClassification and aggregate_diagnostic_superclasses in pyhealth/tasks/ptbxl.py. Non-blocking: a duplicate PTBXLDataset import in pyhealth/datasets/init.py:55 and a duplicate toctree entry in docs/api/datasets.rst:308, plus no tests for the all, diagnostic, subdiagnostic and form label sets. Tests cover the loss terms, shapes, stage switching and projection on synthetic tensors, but nothing exercises the example or the feature_extractor BatchNorm behavior. Overall risk is medium-high as submitted: the example crash and the BatchNorm bug are small fixes, but the duplicated helpers should be consolidated before merge.
Reviewed commit b6d1ee8
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a multi-label PTB-XL task and a ProtoECGNet prototype-based model, with docs, an example and a large synthetic-data test suite. The model is well covered (97%) and the repository total (72%) meets the bar. It is close, but a few problems block merging.
To get this merged:
- The example script crashes.
split_by_strat_fold(samples, folds=[range(1, 9), [9], [10]])does not match the splitter's signature, which expects one fold integer per sample. As written it should fail on every run, including--synthetic. Please use the splitter correctly, or split by thestrat_foldvalue in each sample. Then run--syntheticto confirm it works. - BatchNorm stays frozen in the
feature_extractorstage.set_stagemarks a module as frozen when all its direct parameters haverequires_grad=False.PrototypeBranchownsprototype_vectors, which is frozen in that stage, so the whole branch is put in eval mode ontrain(). The trainable backbone's BatchNorm running stats then never update. Please restrict the frozen set to the modules that are actually frozen, for example leaf modules, or exclude modules that have trainable children. Add a test covering this stage. pyhealth/tasks/ptbxl_multilabel_classification.pyis only 68% covered, below the 80% bar. Please add tests for theall,diagnostic,subdiagnosticandformlabel sets. The fallback WFDB reader's error paths (missing header, unsupported format) also need tests.- The task duplicates existing code.
read_wfdb_signaloverlapsload_ptbxl_recordinpyhealth/datasets/ptbxl.py. The superdiagnostic aggregation overlapsPTBXLSuperclassClassificationandaggregate_diagnostic_superclassesinpyhealth/tasks/ptbxl.py. Please reuse the existing helpers where possible. If a separate implementation is needed, say why in the PR.
Optional suggestions:
- Remove the duplicate
PTBXLDatasetimport inpyhealth/datasets/__init__.py. - Remove the duplicate PTBXLDataset toctree entry in
docs/api/datasets.rst.
Thank you for all the effort you've put into this PR! Given its complexity, a maintainer will take over the review from here, and they can ask for more AI reviews whenever they'd be helpful.
Reviewed commit b6d1ee8
1. examples: split_by_strat_fold(samples) — removed incorrect folds=[range(1,9),[9],[10]]
argument that crashed at runtime; the default 1-8/9/10 split is correct.
2. models/protoecgnet: _frozen_modules now checks recurse=True so a PrototypeBranch
whose backbone has trainable children is never put in eval() during the
feature_extractor stage (backbone BatchNorm was not updating its running stats).
Added test_backbone_bn_stays_in_train_mode_in_feature_extractor_stage to confirm.
3. tests: added 8 tests to TestPTBXLMultilabelClassification covering the all,
diagnostic, subdiagnostic, form label sets, the missing-header error path,
standardize=None, and strat_fold field presence. Task coverage now >80%.
4. tasks/ptbxl_multilabel_classification: read_wfdb_signal delegates to
load_ptbxl_record (master's canonical reader) when wfdb is installed, falling
back to the built-in 16-bit reader otherwise — eliminates duplicate logic for
the wfdb-available path. _aggregate_labels docstring notes it extends
aggregate_diagnostic_superclasses to all six label sets.
Two optional fixes: removed duplicate PTBXLDataset import in datasets/__init__.py
and duplicate toctree entry in docs/api/datasets.rst.
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
Logiquo
added
status: automatic
- examples: add --device {auto,cpu,cuda,mps} flag (default auto picks
CUDA > MPS > CPU) so MacBook CPU-only results are reproducible with
--device cpu instead of being accidentally overridden by MPS.
- examples: seed random and numpy alongside torch at the start of
run_pipeline so the 8-cell ablation grid is reproducible, not just
the first run.
- tests: add test_state_dict_round_trip_preserves_forward_output —
nudges parameters off initializer, saves state_dict into a fresh
model, and asserts forward outputs match exactly. Covers prototypes,
classifier/fusion weights, and buffers.
Not fixed in this commit (documented instead): projection_info is still
not part of state_dict. The "every prediction explained by example"
claim is softened in the PR description and the project_prototypes
docstring (already) explains the save-separately workflow. The paper
--paper-groups path remains unbenchmarked on real PTB-XL; the PR
description now flags that explicitly.
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. The PTB-XL multi-label task, the ProtoECGNet prototype model, the example/ablation script, the docs and the extensive tests are a valuable addition. The model code looks well structured and its coverage (97%) is strong. It is close, but a few things need fixing first.
To get this merged:
- Fix the
state_dictloading problem.class_weightsandlabel_cooccurrenceare registered asNonebuffers (register_buffer(..., None)) and are later assigned tensors. A checkpoint saved from a model built with these statistics will then contain keys that a freshly constructed model lacks, so strictload_state_dictfails. Either register them as real tensors (for example zeros or ones of the right shape) or store them as persistent buffers whose shape is known up front. Also extendtest_state_dict_round_trip_preserves_forward_outputto setclass_weightsandlabel_cooccurrence, since its docstring says it covers them but it never sets them. - Reduce duplication with existing PTB-XL code.
_aggregate_labelsoverlaps withaggregate_diagnostic_superclasses. The fallback 16-bit WFDB reader inread_wfdb_signaloverlaps withload_ptbxl_record. Please reuse the existing helpers where possible, or extend them. Alternatively, explain in the docstring and PR why a separate implementation is needed, and keep the fallback reader only if it is really required. - Correct the interpretability claim. The module docstring and PR description say every prototype is snapped to a real ECG segment. In global mode (
proto_time_len=None) the prototype is a time-averaged feature vector, not a raw segment. Please say so in the module docstring and the PR description. Theproject_prototypesdocstring already notes this. - Strengthen the tests.
test_standardize_none_keeps_raw_scaleonly assertsassertIsNotNone. Please assert actual behaviour, for example that the standardized output has per-lead mean near 0 and std near 1 while the raw output does not. Add a check thatremove_baseline=Truechanges the signal. Coverage ofptbxl_multilabel_classification.pyis 72.9%, below the 80% bar, and these tests plus one for the fallback reader branch would help. The whole-repo coverage of 72% meets the 60% bar.
Optional suggestions:
set_pos_weightis inherited but has no effect on this model. This is documented, but you could override it to raise a warning.- The
--syntheticexample path is not exercised by any test. A small smoke test would help. docs/api/models.rstends with a stray blank line.
Reviewed commit ca0f9a3
claim, test coverage and dedup justifications
- protoecgnet: register class_weights and label_cooccurrence as
real-tensor buffers (ones/zeros with known shape) instead of None.
Previously a checkpoint saved from a model with label statistics
carried keys that a freshly constructed model lacked, so strict
load_state_dict failed. Forward uses label_cooccurrence.any() to
gate the contrastive term so defaulting to zeros is a no-op.
- test_state_dict_round_trip_preserves_forward_output now actually
calls set_label_statistics before saving, so the round-trip covers
the two buffers its docstring claims to cover.
- protoecgnet module docstring: correct the interpretability claim.
Projected prototypes are feature-space representations; in global
mode (proto_time_len=None) specifically, the stored vector is a
time-averaged feature rather than a raw segment. Case-based
explanations should display the matching window from the source
sample, retrieved via the returned projection mapping.
- ptbxl_multilabel_classification: explain in the docstrings why
read_wfdb_signal keeps a fallback (so tests + --synthetic run
without the optional wfdb extra) and why _aggregate_labels doesn't
call through to aggregate_diagnostic_superclasses (the public
helper doesn't take a label set or likelihood threshold).
- tests: replace the weak test_standardize_none_keeps_raw_scale
(asserted assertIsNotNone) with test_standardize_record_versus_none
that checks per-lead mean≈0 / std≈1 for 'record' and non-zero mean
for None. Add test_remove_baseline_changes_signal and three
wfdb-import-blocked tests (unsupported format, missing .dat,
successful roundtrip) that exercise the fallback reader branches
CI's wfdb install would otherwise skip. Task coverage expected to
cross 80% once CI re-runs.
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a multi-label PTB-XL task with six SCP label sets, a ProtoECGNet re-implementation with staged training and prototype projection, an example/ablation script, docs and 59 tests. Coverage is good: protoecgnet.py is at 97.5% and the task file at 96.3%, both above the 80% bar, and the whole repository is at 72.2%, above the 60% bar. The model math looks sound.
To get this merged:
- Duplication with the existing superclass task.
PTBXLMultilabelClassificationrepeats the record loop and the superclass aggregation of the existingPTBXLSuperclassClassificationinpyhealth/tasks/ptbxl.py. Itssuperdiagnosticoutput matches that task. Please reuse or extend the existing code instead of copying it. For example, share the record loop and signal loading, and add only what is new: the other label sets, likelihood filtering, and the baseline and standardization options. - Second WFDB decoder.
read_wfdb_signaladds a fallback decoder next toload_ptbxl_record. As a result, the decoded signal depends on whetherwfdbis installed. Please useload_ptbxl_recordonly. In the tests, mock it or skip whenwfdbis missing, so there is one decoding path. set_pos_weightis silently ignored.ProtoECGNetinheritsBaseModel.set_pos_weightbut does nothing with it, and the code that calls it (for examplepos_weight="balanced"in autoimprove) would be dropped without any error. Please either honor it by mapping it toclass_weights, or override it to raise a clear error or warning.
Optional suggestions:
- Revert the unrelated title change in
docs/api/datasets.rst. - Consider moving the example under
examples/ecg/ptbxl/. - Use
int(float(event.strat_fold)), as the existing task does, to be more tolerant of the CSV format. - The reported results use a single seed and show no consistent gain from the contrastive loss. A short note on that in the PR would help readers.
Reviewed commit be90e67
claim, test rigor and docstring justifications
The round-trip test added in the last commit passed only because
class_weights and label_cooccurrence were never actually set — the
underlying buffers were registered with register_buffer(name, None)
and promoted to tensors only when supplied, so a checkpoint saved
from a model with real label statistics carried keys a freshly
constructed model lacked, and strict load_state_dict failed. Four
fixes in response:
- protoecgnet: register class_weights as torch.ones(num_labels) and
label_cooccurrence as torch.zeros(num_labels, num_labels). ones is
BCE's pos_weight identity, so the default is a no-op. The forward
pass gates the contrastive term with
bool(self.label_cooccurrence.any()), so an all-zeros buffer leaves
the term off; the check is hoisted out of the per-branch loop so
it fires once per forward pass. set_label_statistics docstring now
spells out the buffer semantics.
- test_state_dict_round_trip_preserves_forward_output: call
set_label_statistics(class_weights=..., label_cooccurrence=...)
before the save, so the test actually exercises the buffers its
docstring claims to cover; also assert the two buffers match after
load.
- protoecgnet module docstring: correct the interpretability claim.
Projected prototypes are feature-space representations, and in
global mode (proto_time_len=None) the stored vector is a
time-averaged feature rather than a raw ECG segment. Case-based
explanations should display the matching window from the source
sample, retrieved via the returned source-sample + window mapping.
PR description updated to match.
- ptbxl_multilabel_classification: strengthen the docstrings
justifying read_wfdb_signal's fallback reader (needed so tests and
--synthetic run without the optional wfdb extra; the wfdb-backed
path still delegates to load_ptbxl_record) and _aggregate_labels
(the public aggregate_diagnostic_superclasses doesn't take a label
set or a likelihood threshold, so a call-through isn't viable).
- tests: replace the weak test_standardize_none_keeps_raw_scale
(asserted only assertIsNotNone) with test_standardize_record_versus_none
that checks per-lead mean≈0 / std≈1 under 'record' and non-zero
per-lead means under None. Add test_remove_baseline_changes_signal
(filtered output differs, and its per-lead mean is closer to zero).
Add three fallback-reader tests that block the wfdb import to force
the fallback branches CI's wfdb install would otherwise skip
(unsupported format, missing .dat, successful roundtrip). Task
coverage expected to clear the 80% bar once CI reruns.
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
honor set_pos_weight, plus the four optional suggestions
Three blocking issues:
1. Second WFDB decoder removed. read_wfdb_signal is deleted, the
fallback 16-bit reader goes away, and _prepare_signal calls
pyhealth.datasets.ptbxl.load_ptbxl_record directly — the same
decoding path PTBXLSuperclassClassification uses. There is now one
decoding path, so the decoded signal no longer depends on install
state. test_ptbxl.py adds a module-level _WFDB_AVAILABLE check and
@skipUnless on TestPTBXLMultilabelClassification; the three
builtins.__import__-blocking fallback tests are removed because
the code they covered is gone. A new
test_load_signal_matches_superclass_task_signal asserts the
multi-label task's _prepare_signal(standardize=None) matches
load_ptbxl_record exactly.
2. Reuse PTBXLSuperclassClassification's aggregation. For
label_set="superdiagnostic" with min_likelihood=0 the task's
__call__ now routes through aggregate_diagnostic_superclasses +
load_diagnostic_class_map, so the labels produced for the
overlapping configuration match PTBXLSuperclassClassification
record-for-record. Other label sets or non-zero likelihood
thresholds keep using the generalising _aggregate_labels helper
(which the public helper cannot replace — it doesn't take a label
set or a likelihood threshold). Guarded by
test_superdiagnostic_matches_superclass_task.
3. set_pos_weight no longer silently ignored. ProtoECGNet now
overrides BaseModel.set_pos_weight to delegate to the base
implementation (parses scalars, sequences, tensors and "balanced")
and then mirrors self._pos_weight into self.class_weights. None
resets class_weights to the ones default. The class docstring Note
is corrected to describe the routing instead of a no-op. Guarded
by test_set_pos_weight_routes_to_class_weights.
Four optional suggestions:
- docs/api/datasets.rst: revert the toctree title — master's titled
entry (PTB-XL Dataset <...>) is restored and the untitled entry
introduced on this branch is removed.
- Example moved from examples/ to
examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py
(next to ptbxl_superclass_quickstart.py). PR description and run
commands updated to the new path.
- strat_fold now reads as int(float(event.strat_fold)) to match
PTBXLSuperclassClassification and tolerate "9.0" strings from the
CSV.
- PR description now says results use one seed per cell and that the
contrastive loss did not show a consistent gain over
--no-contrastive at this compute budget; a multi-seed sweep would
be needed to disentangle signal from seed noise.
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a multi-label PTB-XL task, the ProtoECGNet prototype model, an end-to-end example with an ablation grid, and docs. The tests are thorough, and coverage clears the bar. The whole repository is at 72.16% against the 60% bar. Every added or modified file under pyhealth/ is above 80%, apart from pyhealth/datasets/init.py at 83.33%, which is also above it.
To get this merged:
- Fix the shape of
class_weightsafterProtoECGNet.set_pos_weight. When the argument is a scalar, the base-class parsing gives a (1,) tensor, and your override assigns it directly to theclass_weightsbuffer. That changes the buffer's shape from (num_labels,) to (1,). A checkpoint saved afterset_pos_weight(4.0)will then fail a strictload_state_dictinto a freshly built model because the shapes differ. Expand the parsed weight to (num_labels,) in the override, for example.expand(self.num_labels).clone()when it has a single element. Also check that a wrong-length vector raises a clear error. - Update
test_set_pos_weight_routes_to_class_weights. It currently assertstorch.tensor([4.0])for the scalar case. Change it to expect a length-num_labels tensor filled with 4.0, and ideally add a save/load round trip afterset_pos_weight.
Optional suggestions:
- Add a test for
set_pos_weight("balanced", dataset), which is currently uncovered. - The usage lines in the example docstring point to
examples/ptbxl_multilabel_classification_protoecgnet.py, but the file lives atexamples/ecg/ptbxl/.... - A Trainer-level test would help cover the staged training flow beyond the --synthetic example run.
The results are single-seed. You state that clearly, and that is fine as it stands.
Reviewed commit f290e48
set_pos_weight, balanced test, example path, trainer-level test
Blocker:
set_pos_weight(scalar) left class_weights at shape (1,) because the
base class's reshape(-1) produces a (1,)-tensor for a scalar. A
subsequent state_dict() carried the (1,) shape, which failed strict
load_state_dict into a freshly constructed model with the (num_labels,)
default. Fix in protoecgnet.set_pos_weight: when the parsed weight has
numel()==1, expand to (num_labels,) and clone; when it has
numel()==num_labels, assign as-is; otherwise raise ValueError with a
clear message. Docstring updated to describe the broadcast and the
error case.
Test coverage:
- test_set_pos_weight_routes_to_class_weights now asserts
class_weights.shape == (num_labels,) for both the scalar and the
per-label paths, exercises the wrong-length ValueError, and ends
with a state_dict round-trip after set_pos_weight(2.5) that
verifies strict load_state_dict works.
- Add test_set_pos_weight_balanced_uses_dataset: calls
set_pos_weight("balanced", dataset), checks the resulting shape and
that all weights are finite and positive, and checks that
"balanced" without a dataset raises.
- Add test_trainer_runs_through_staged_protocol: a real
pyhealth.trainer.Trainer walks through feature_extractor → joint →
project_prototypes → classifier, each for 1 epoch, and the final
inference call must return finite logits across all labels.
Optional:
- examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py
usage docstring: all three CLI examples now point at the file's
current path (examples/ecg/ptbxl/...), not the old examples/ root.
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
There was a problem hiding this comment.
Previous review
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a multi-label PTB-XL task, a ProtoECGNet implementation, an example with an ablation script, docs, and a large synthetic-data test suite. The work is thorough, and you state its limits openly. Coverage is above the bar: protoecgnet.py is at 97.6%, the new task file at 97.5%, and the whole repository at 73.7% against the 60% minimum. One correctness problem needs fixing first.
To get this merged:
- Persist the training stage with the checkpoint.
ProtoECGNet.stageis a plain attribute and is not instate_dict. A model built withfusion="learned"and then reloaded starts instage="all". Itsforwardtherefore returns the branch-scatter logits instead of the fusion head output.prototype_contributionslikewise returns branch-classifier contributions instead of fusion ones. Nothing warns the user. You could store the stage viaget_extra_state/set_extra_state, or register it as a buffer. Alternatively, raise or warn when a learned-fusion model is used without an explicitset_stage. If you keep the current behaviour, document it prominently in the class docstring. - Cover this in a test.
test_state_dict_round_trip_preserves_forward_outputusesfusion="learned"but never callsset_stage("fusion"). It compares outputs in the default stage, so it can't catch this problem. Please add a case that reloads a model in thefusionstage and checks that the logits match.
Optional suggestions:
_aggregate_labelsre-implements part of the superclass aggregation logic. That's reasonable given the extra label sets and the likelihood filter, but a short note in the code on why it can't share more would help future maintainers.- The reported results are single-seed. You already say so, and that's fine.
Reviewed commit 2216e3a
checkpoints reload with the correct forward path
Blocker:
self.stage was a plain Python attribute, not part of state_dict. A
ProtoECGNet(..., fusion="learned") trained and saved in the "fusion"
stage reloaded into stage="all", so forward returned branch-scatter
logits instead of fusion_classifier output and prototype_contributions
returned branch-classifier contributions instead of fusion ones.
Silent, no warning. Fix with the PyTorch-native hook pair:
- ProtoECGNet.get_extra_state() returns {"stage": self.stage}; PyTorch
stores this under "_extra_state" in state_dict.
- ProtoECGNet.set_extra_state(state) calls self.set_stage(state["stage"])
on load, which also recomputes _frozen_modules so BN freezing stays
correct after reload.
Test coverage:
The existing test_state_dict_round_trip_preserves_forward_output used
fusion="learned" but never called set_stage("fusion"), so it couldn't
catch this. Add test_state_dict_round_trip_preserves_fusion_stage: a
model is moved into the fusion stage, parameters are nudged, then it
is saved and loaded into a fresh model. Asserts clone.stage == "fusion"
after load, logits match exactly, and prototype_contributions match
per-branch (also pins the fusion-branch contribution path).
Signed-off-by: Irvi Aini <irvi.fa@gmail.com>
|
/review |
1 similar comment
|
/review |
| outputs = {n: self.branches[n](x) for n in self.branch_names} | ||
| sims = torch.cat([o[1] for o in outputs.values()], dim=1) | ||
| logits = self._scatter({n: o[0] for n, o in outputs.items()}, batch) | ||
| if self.stage == "fusion": |
There was a problem hiding this comment.
[medium] Learned fusion head is ignored outside the 'fusion' stage
In forward(), logits come from the branch-scatter path unless self.stage == "fusion". With fusion="learned", the default stage "all" and the "joint"/"classifier" stages therefore never use or train fusion_classifier, and it stays at its random init. The class docstring says "all" trains everything with one loss, and the PR describes learned fusion as supported, so this does not match the description. test_two_branches_with_learned_fusion only checks output shapes in the default stage, so the gap is not caught. Fix: use the fusion head whenever it exists in stages that compute the loss (or raise in "all" when fusion='learned'), and add a test that default-stage logits depend on fusion_classifier.
| """ | ||
| y = labels.float() | ||
| counts = y.sum(dim=0).clamp_min(1.0) | ||
| weights = (counts.sum() - counts) / counts |
There was a problem hiding this comment.
[low] class_weights use a different total than BaseModel's 'balanced' option
compute_label_statistics uses counts.sum() (the sum of label counts) as the total for (total - count) / count. BaseModel.set_pos_weight('balanced') (pyhealth/models/base_model.py, around line 257) uses the sample count. Both write to the same class_weights buffer, so its meaning depends on the path used. The docstring cites the reference code, which I could not verify. Fix: use the sample count, or document the difference and where it comes from.
| signal = (signal - mean) / (std + 1e-6) | ||
| return np.ascontiguousarray(signal, dtype=np.float32) | ||
|
|
||
| def __call__(self, patient: Any) -> list[dict[str, Any]]: |
There was a problem hiding this comment.
[low] New task duplicates the record loop of PTBXLSuperclassClassification
call repeats the per-record loop and label plumbing of PTBXLSuperclassClassification (pyhealth/tasks/ptbxl.py, lines 155-270). The superdiagnostic path already delegates to the same helpers, and the superclass docstring says the 71-SCP work belongs in that module. The new task is a reasonable extension (label sets, likelihood filter, preprocessing), so this is not redundant. It could instead be added as parameters to the existing class, or the shared loop could be factored out.
There was a problem hiding this comment.
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a multi-label PTB-XL task and a prototype-based ProtoECGNet model, with an example, an ablation script and docs, and the tests are thorough. Coverage meets the bar: protoecgnet.py is at 97.6% and ptbxl_multilabel_classification.py at 97.5%, and the repository total is 73.7% against the 60% minimum. It looks close.
To get this merged:
- Learned fusion outside the "fusion" stage. In ProtoECGNet.forward, fusion_classifier is used only when stage == "fusion". With fusion="learned" and the default "all" stage, or "joint" or "classifier", the fusion layer is never used or trained. The logits come from the per-branch scatter, so the model behaves like fusion="average" with extra unused parameters. Please make this explicit in one of these ways: raise an error or warning when fusion="learned" is used in a stage that bypasses the fusion head; route the logits through the fusion head in the "all" stage; or state the behaviour clearly in the fusion and set_stage docstrings.
- Test for that behaviour. test_two_branches_with_learned_fusion checks only output shapes. Please add an assertion that covers the chosen behaviour, for example that the fusion head affects the logits only in the "fusion" stage.
Optional suggestions:
- compute_label_statistics computes class weights as (sum of label counts − count) / count, which differs from the "balanced" meaning in BaseModel.set_pos_weight. A short docstring note would help users who mix the two.
- The call record loop closely mirrors PTBXLSuperclassClassification. A shared helper would reduce duplication later.
- _aggregate_labels is documented as a necessary reimplementation, which is reasonable. Consider extending the shared helper later.
Thank you for all the effort you've put into this PR! Given its complexity, a maintainer will take over the review from here, and they can ask for more AI reviews whenever they'd be helpful.
Reviewed commit 968f2cb
github-actions
Bot
added
status: manual
document class_weights difference from BaseModel.set_pos_weight
Blocker 1 (medium): learned fusion head ignored outside "fusion" stage.
With fusion="learned" and self.stage == "all" (the default, and the
stage the class docstring describes as "trains everything with one
loss"), forward() went through the branch-scatter path and never
touched fusion_classifier. The fusion head stayed at random init and
never received gradients, so the model silently behaved like
fusion="average" with extra unused parameters. Same mismatch in
prototype_contributions.
Fix: add ProtoECGNet._uses_fusion_head() returning True iff
fusion_classifier is not None and stage in {"all", "fusion"}. forward()
and prototype_contributions() both route through fusion_classifier
when this is True, keeping the paper's staged protocol intact for
"joint" / "classifier" / "prototypes" (which still go through the
per-branch classifiers so the branches get trained in isolation
before the fusion head is fine-tuned). The fusion argument's docstring
and prototype_contributions's Returns block are updated to point at
_uses_fusion_head.
Tests:
- test_fusion_head_affects_logits_in_all_and_fusion_stages: perturb
fusion_classifier.weight and check that logits change for stage in
{"all", "fusion"} and are unchanged for stage in
{"joint", "classifier"}. Catches regressions in both directions.
- test_all_stage_trains_fusion_head_with_learned_fusion: backprop
through model(**batch)["loss"] in the "all" stage must give
fusion_classifier.weight a non-zero gradient.
Blocker 2 (low): class_weights total differs from BaseModel.set_pos_weight.
compute_label_statistics uses counts.sum() (sum of per-label positive
counts) as the normalization total, following the paper's reference
code. BaseModel.set_pos_weight("balanced") uses the sample count
instead, and both end up in the same class_weights buffer. Docstring
Note added spelling out the difference and recommending callers mix
them deliberately rather than by accident.
|
@jhnwu3 All prior blockers addressed. Would you review, or route to someone who can? |
# Conflicts: # docs/api/tasks.rst # pyhealth/tasks/__init__.py
Contributor: NetID
iaini2(iani2@illinois.edu)Course: CS 598 DLH (Fall 2026), PyHealth reproduction project, Option 4 (Task + Model)
Summary
PTBXLMultilabelClassification(pyhealth/tasks/ptbxl_multilabel_classification.py):Converts each labelled ECG from
PTBXLDatasetinto a(12, 10 * fs)float32 signalwith multi-hot labels. Supports all six SCP-ECG label sets (
all,diagnostic,subdiagnostic,superdiagnostic,form,rhythm), 100/500 Hz, likelihood filtering,per-record z-score normalisation and an optional 0.5 Hz high-pass baseline removal.
Shares the single
wfdb-backed decoding path withPTBXLSuperclassClassification(
load_ptbxl_recordfrompyhealth.datasets.ptbxl); forlabel_set="superdiagnostic"with
min_likelihood=0the label aggregation delegates toaggregate_diagnostic_superclasses, so the labels produced by both tasks matchrecord-for-record (guarded by
test_superdiagnostic_matches_superclass_task).ProtoECGNet(pyhealth/models/protoecgnet.py):Re-implementation of Sethi et al., MLHC 2025 (https://arxiv.org/abs/2504.08713).
A 1D or 2D ResNet extracts features; each prototype branch computes cosine or L2
similarities between feature windows and learned prototype vectors assigned to labels;
a final linear layer produces multi-label logits. Why prototype-based? Standard
ECG classifiers are black boxes. After calling
project_prototypes(), every prototypeis replaced by the feature-space representation of a real training ECG segment (in
global mode,
proto_time_len=None, it is a time-averaged feature rather than a rawsegment); the method returns the source-sample and window indices so that the matching
ECG window can be retrieved for case-based explanations.
prototype_contributions()then returns per-prototype contributions to each logit,which combined with the projection mapping (saved alongside the checkpoint — see the
method docstring) lets a clinician inspect which training cases drove a prediction.
set_pos_weightis overridden to route the inheritedBaseModelargument intoclass_weights(including the"balanced"option), so callers likeautoimprovework unchanged.
Supports the paper's rhythm / morphology / global branches (16 / 52 / 3 labels),
learned fusion, contrastive and diversity regularisation, and the staged training
protocol (
feature_extractor→joint→ projection →classifier→fusion)through
set_stage()andproject_prototypes().Example (
examples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.py):End-to-end pipeline + ablation (contrastive on/off × cosine/L2 × prototype count).
--syntheticruns on synthetic data with no download required.--device {auto,cpu,cuda,mps}controls the training device (defaultautopicksCUDA > MPS > CPU).
Docs: RST pages wired into the existing API reference.
Dataset and how to run
Dataset: PTB-XL v1.0.3 is open access on PhysioNet (no credentialing). Download with:
--rootmust point to the folder containingptbxl_database.csv,scp_statements.csvand
records100/. The task requires the optionalwfdbextra:pip install 'pyhealth[ptbxl]'.Quick check (no download):
Real data:
Paper
Sethi et al., ProtoECGNet: Case-Based Interpretable Deep Learning for Multi-Label ECG
Classification with Contrastive Learning, MLHC 2025 (PMLR 298).
https://proceedings.mlr.press/v298/sethi25a.html
Authors' reference code: https://github.com/bbj-lab/protoecgnet
Differences from the paper and authors' code
BaseModel/Trainer(reference code requires Lightningand hard-coded cluster paths).
PTBXL_BRANCH_KWARGS; defaultis
--prototypes-per-labelsince the paper's counts are heavy for a free GPU.feature_extractorstage pretrains each branch on its own labels; the paper pretrainsblack-box models on all 71.
classifier/fusionstages (--lam-l1, off by default).(
proto_time_len=None) specifically, the stored vector is a time-averaged feature.projection_info(the per-prototype source-sample + window mapping produced byproject_prototypes) is not part ofstate_dictand must be saved separatelyalongside the checkpoint — see the method's docstring.
Files to review
pyhealth/tasks/ptbxl_multilabel_classification.pyaggregate_diagnostic_superclassesfor the superdiagnostic path)pyhealth/models/protoecgnet.pyset_pos_weightoverridepyhealth/{tasks,models}/__init__.pytests/core/test_ptbxl.py(appended class)tests/core/test_protoecgnet.pyexamples/ecg/ptbxl/ptbxl_multilabel_classification_protoecgnet.pydocs/api/{tasks,models}.rstand two new.rstpagesTests
pytest tests/core/test_ptbxl.py tests/core/test_protoecgnet.py: 56 tests, all onsynthetic data. No real data is read. The
TestPTBXLMultilabelClassificationsuite is@skipUnless(wfdb_installed). Cover loss terms (per-sample clustering andseparation, Eqs. 4–5), prototype shapes per branch, projection onto real segments, stage
switches, label grouping, all six label sets, a BN-stats regression for the
feature_extractorfreeze fix, astate_dictround-trip (withclass_weightsandlabel_cooccurrenceset) that checks forward outputs match after save/load,behavioural assertions for
standardize='record'(per-lead mean≈0 / std≈1) andremove_baseline=True, record-for-record parity withPTBXLSuperclassClassificationon the superdiagnostic path, exact signal match withload_ptbxl_record, andset_pos_weightrouting toclass_weights.Results
Real PTB-XL runs on the 5
superdiagnosticlabels (single branch, 1D ResNet-18,3 prototypes per label), folds 1–8 / 9 / 10 (train 17,084 / val 2,146 / test 2,158),
Adam 1e-3, batch 32, one seed per cell, invoked with
--device cpuor on aKaggle T4.
Note that the contrastive loss (
--lam-cnrst 300, the paper's default) does not showa consistent gain over
--no-contrastivein these single-seed runs. A multi-seedsweep would be needed to distinguish signal from seed noise; the paper reports an
Optuna search of hundreds of trials per variant, well beyond this reproduction's
compute budget.
Compute
Local: MacBook Pro M4 Pro (12 cores, 24 GB), CPU only (
--device cpu).Remote: Kaggle Tesla T4 (16 GB), 4 vCPU Xeon 2.0 GHz, 31 GB RAM.