diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index fdf0780..2df946a 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -3,14 +3,7 @@ .gitmodules @StevenSong .pre-commit-config.yaml @StevenSong .secrets.baseline @StevenSong -configs/audio @StevenSong -configs/multi-arm-heedb @StevenSong -configs/pretrain-supervised.yaml @StevenSong -configs/pretrain-unsupervised.yaml @StevenSong -configs/target-blackbox.yaml @StevenSong -configs/target-guided-14ppl.yaml @StevenSong -configs/target-guided-5ppl.yaml @StevenSong -configs/target-unguided.yaml @StevenSong +configs/ @StevenSong data-preprocessing/ @StevenSong env.yaml @StevenSong external/PierreElias-IntroECG @StevenSong @@ -41,7 +34,9 @@ protossl/lightning_utils.py @StevenSong protossl/models/__init__.py @StevenSong protossl/models/_base_classifier.py @StevenSong protossl/models/_blackbox_classifier.py @StevenSong +protossl/models/_blackbox_contraster.py @StevenSong protossl/models/_pretrained_utils.py @StevenSong +protossl/models/_protoecgnet.py @StevenSong protossl/models/_prototype_assigner.py @StevenSong protossl/models/_prototype_classifier.py @StevenSong protossl/models/_prototype_contraster.py @StevenSong @@ -60,7 +55,9 @@ protossl/models/helpers/_prototype_ilp_assigner.py @sahilsethi0105 protossl/models/layers/__init__.py @StevenSong protossl/models/layers/_multi_input_linear.py @StevenSong protossl/models/layers/_panns_backbones.py @StevenSong +protossl/plotting.py @StevenSong protossl/trainer.py @StevenSong +protossl/protoecgnet_trainer.py @StevenSong pyproject.toml @StevenSong README.md @StevenSong @sahilsethi0105 requirements.txt @StevenSong @@ -72,40 +69,33 @@ scripts/_slurm_wrapper.sh @StevenSong scripts/_submit_job.sh @StevenSong scripts/0-run-cache-data.sh @StevenSong scripts/1-run-blackbox-direct.sh @StevenSong -scripts/2-run-labsup-proto-direct.sh @StevenSong +scripts/2-run-supproto-direct.sh @StevenSong scripts/2-z-* @StevenSong -scripts/3-run-protossl-heedb-pila.sh @StevenSong -scripts/3-z-* @StevenSong -scripts/4-run-labsup-proto-heedb-rila.sh @StevenSong -scripts/4-y-* @StevenSong -scripts/4-z-* @StevenSong +scripts/3-* @StevenSong +scripts/4-* @StevenSong scripts/5-1-run-ecgfounder-logreg.sh @StevenSong scripts/5-2-run-stmem-logreg.sh @StevenSong +scripts/5-3-run-pclr-logreg.sh @StevenSong scripts/6-run-protossl-heedb-pia.sh @StevenSong scripts/6-z-* @StevenSong -scripts/7-run-protossl-heedb-pit.sh @StevenSong -scripts/8-run-protossl-heedb-pip.sh @StevenSong +scripts/7-* @StevenSong scripts/9-0-run-ecgfounder-patches.sh @StevenSong scripts/9-1-run-ecgfounder-lap.sh @StevenSong scripts/9-2-run-ecgfounder-clustering.sh @StevenSong scripts/9-3-run-ecgfounder-random.sh @StevenSong scripts/audio @StevenSong -scripts/ecg-fms/_compute_ecgfounder_embeddings.py @StevenSong -scripts/ecg-fms/_compute_stmem_embeddings.py @StevenSong +scripts/ecg-fms/ @StevenSong scripts/echonext/_tabular_logreg.py @StevenSong scripts/echonext/run-columbia-minimodel.sh @StevenSong scripts/echonext/run-tabular-logreg.sh @StevenSong -scripts/pretrain/_cache_heedb.py @StevenSong -scripts/pretrain/_compute_heedb_normalizations.py @StevenSong -scripts/pretrain/run-heedb-normalizations.sh @StevenSong -scripts/pretrain/run-pass-heedb-pretrain-no-attn.sh @StevenSong -scripts/pretrain/run-pass-heedb-pretrain.sh @StevenSong -scripts/pretrain/run-prosup-heedb-pretrain.sh @StevenSong +scripts/pretrain/ @StevenSong scripts/prototypes-from-fms @StevenSong scripts/queue-experiments.sh @StevenSong +scripts/queue-pclr-logreg.sh @StevenSong scripts/README.md @StevenSong @sahilsethi0105 user-study/images @StevenSong user-study/metadata.csv @StevenSong user-study/prepare_samples.ipynb @StevenSong user-study/results.csv @StevenSong user-study/user_study_analysis.ipynb @sahilsethi0105 +large-user-study/ @StevenSong \ No newline at end of file diff --git a/.gitignore b/.gitignore index 2c32da1..c21884a 100644 --- a/.gitignore +++ b/.gitignore @@ -219,3 +219,6 @@ OLD/ slurm-logs/ results/figs/ ecgfounder-checkpoint/ +wandb/ +.devcontainer/ +user-study/images-pdf diff --git a/README.md b/README.md index 8b795be..8b9f929 100644 --- a/README.md +++ b/README.md @@ -1,20 +1,34 @@ -# ProtoSSL: Interpretable Prototype Learning from Unlabeled Time-Series Data +# ProtoSSL: Self-Supervised Pretraining and Downstream Transfer for Projection-Based Prototype Models -This repo contains the source code for **ProtoSSL**, a novel framework for label-free learning of interpretable, projection-based prototypes that are readily adaptable to downstream tasks. Our key innovation is to separate *motif discovery* from *label alignment*. ProtoSSL first learns a reusable prototype bank using a self-supervised objective applied directly to prototype activations, and then aligns these prototypes to downstream tasks through a novel and efficient assignment procedure. We study ProtoSSL and provide code for our experiments primarily over ECGs and additionally audio waveforms. +This repo contains the source code for **ProtoSSL**, a novel framework for pretraining a foundational latent prototype bank on unlabeled data and transferring it to downstream tasks to create interpretable, projection-based prototype models. Our key innovation is to separate *motif discovery* from *label alignment*. ProtoSSL first learns a reusable latent prototype bank using a self-supervised objective applied directly to prototype activations, and then aligns these prototypes to downstream tasks through a novel and efficient assignment procedure. We study ProtoSSL and provide code for our experiments primarily over ECGs and additionally audio waveforms. ### Citation ``` -TODO: bibtex citation +@article{song2026protossl, + title={ProtoSSL: Self-Supervised Pretraining and Downstream Transfer for Projection-Based Prototype Models}, + author={Song, Steven and Sethi, Sahil and Beaulieu-Jones, Brett and Grossman, Robert L}, + journal={arXiv preprint arXiv:2605.06943}, + year={2026} +} ``` -## Contributing +## Repo Organization + +| Path | Contents | +|------|----------| +| `protossl/` | The `protossl` package: datasets, models, and the Lightning trainers (`protossl.trainer` for ProtoSSL and baselines, `protossl.protoecgnet_trainer` for multi-arm ProtoECGNet) | +| `configs/` | Lightning CLI configs for pretraining (`pretrain-*.yaml`) and downstream tasks (`target-*.yaml`), with audio configs under `configs/audio/` | +| `scripts/` | Run scripts for every experiment in the paper, plus the slurm queue scripts that submit them. See [`scripts/README.md`](scripts/README.md) | +| `user-study/` | Sample preparation and analysis for the user study (`user-study/images` is a submodule of the de-identified study images) | +| `results/` | Notebooks that produce the figures and tables in the paper from the experiment outputs | +| `data-preprocessing/` | Notebooks that create the train/val/test splits and nested training subsets, plus precomputed HEEDB normalization statistics. See [`data-preprocessing/README.md`](data-preprocessing/README.md) | +| `external/` | Submodules for baselines: [ST-MEM](https://github.com/vuno/ST-MEM) and our fork of [IntroECG](https://github.com/PierreElias/IntroECG) (EchoNext minimodel) | -**Repo Organization:** The outline of this repo is centered around the package `protossl`, defined in the subdirectory of the same name. While the package definition enables experimentation, the `scripts` folder defines all of the actual experiment run scripts. Most of the experimental comparisons we run have standalone implementations also contained within the `scripts` folder. +## Setup -**Environment Setup:** We use pre-commit hooks to maintain coding style and to enforce data stewardship. Source data should never be committed to this repo, including exploratory notebooks which may accidentally leak source data. The pre-commit hooks do not necessarily prevent this and can be circumvented, but may help prevent obvious cases. To correctly set up your environment: ```bash -# 1) clone repo -git clone git@github.com:StevenSong/ProtoSSL.git +# 1) clone repo, including submodules (only needed for the ST-MEM and EchoNext minimodel baselines, and the user study images) +git clone --recurse-submodules git@github.com:StevenSong/ProtoSSL.git cd ProtoSSL # 2) create and activate environment @@ -22,16 +36,17 @@ cd ProtoSSL conda env create -f env.yaml conda activate protossl -# 3) enable pre-commit hooks -pre-commit install - -# 4) install protossl from editable source +# 3) install protossl from editable source pip install -e . -# 5) dev away +# 4) (contributors only) enable pre-commit hooks, see Contributing below +pre-commit install ``` -**torchcodec:** torchcodec is a bit fragile with dependencies. We've pinned `torch==2.7.0` which is compatible with `torchcodec==0.4.0`, both compiled against CUDA 12.8 (which we use on our machines). This torchcodec version is only compatible with `datasets==4.0.0`. If you see errors relating to torchcodec (you can diagnose this by just importing torchcodec), make sure the dependencies are compatible not just relative to versioning, but also relating to the CUDA versions. We also use `ffmpeg=5.*` installed via conda. If you see an error relating to not being able to find `libnppicc.so.12`, it might be that the linker can't find the binaries (which we ensure are available by installing `nvidia-npp-cu12`). To fix this, you can try setting the `LD_LIBRARY_PATH` environment variable: +
+Troubleshooting torchcodec + +torchcodec is a bit fragile with dependencies. We've pinned `torch==2.7.0` which is compatible with `torchcodec==0.4.0`, both compiled against CUDA 12.8 (which we use on our machines). This torchcodec version is only compatible with `datasets==4.0.0`. If you see errors relating to torchcodec (you can diagnose this by just importing torchcodec), make sure the dependencies are compatible not just relative to versioning, but also relating to the CUDA versions. We also use `ffmpeg=5.*` installed via conda. If you see an error relating to not being able to find `libnppicc.so.12`, it might be that the linker can't find the binaries (which we ensure are available by installing `nvidia-npp-cu12`). To fix this, you can try setting the `LD_LIBRARY_PATH` environment variable: ```bash export LD_LIBRARY_PATH=$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/npp/lib:$LD_LIBRARY_PATH # test by importing torchcodec in a python runtime @@ -47,3 +62,49 @@ echo 'export LD_LIBRARY_PATH=$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/n # Unset on deactivate echo 'export LD_LIBRARY_PATH=$(echo $LD_LIBRARY_PATH | sed "s|$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/npp/lib:||g")' > $CONDA_PREFIX/etc/conda/deactivate.d/npp_lib.sh ``` +
+ +## Data + +We do not redistribute any data; each dataset must be obtained from its original source under its own terms of use. + +* **ECG pretraining:** HEEDB +* **ECG downstream:** EchoNext, PTB-XL, CinC 2020 (Georgia), MIMIC-IV-ECG, ZZU pECG, CODE-15% +* **Audio pretraining:** AudioSet +* **Audio downstream:** UrbanSound8K, IEMOCAP, VoxCeleb1 + +The `data-preprocessing/` notebooks create the splits and nested training subsets (e.g. `ptb-xl-1k`) used for our data-efficiency experiments. + +## Reproducing Our Experiments + +1. Obtain the datasets above and create the splits and training subsets with the `data-preprocessing/` notebooks +1. Find and replace the hardcoded paths in the run, queue, and pretrain scripts: `/opt/gpu_working/steven` (repo and outputs) and `/opt/gpudata` (source datasets) +1. Pretrain: run the scripts in [`scripts/pretrain/`](scripts/pretrain/) for ECG and `scripts/audio/run-pretrain-audio-*.sh` for audio +1. Submit the downstream experiments via slurm: [`scripts/queue-experiments.sh`](scripts/queue-experiments.sh) (ECG) and [`scripts/audio/queue-experiments.sh`](scripts/audio/queue-experiments.sh) (audio) +1. Produce figures and tables with the `results/` notebooks + +See [`scripts/README.md`](scripts/README.md) for more detail. Each run script writes to `$RUN_DIR/`, named `--[-suffixes]`: + +* **Models:** `protossl`, `supproto` (supervised prototypes), `blackbox`, and linear probes over foundation model embeddings (`ecgfounder-logreg`, `stmem-logreg`, `pclr-logreg`) +* **Transfer methods:** `pila` (project in linear assignment), `rila` (reproject in linear assignment), `pia` (project in assignment, using ProtoPool-based assignment), `pit` (project in transfer dataset without assignment or label supervision), `pip` (project in pretraining dataset without assignment or label supervision), `latent` (no projection), `direct` (trained only on the target dataset) +* **Suffixes:** `-ft` is fine-tuned, `-Nppl` uses N prototypes per label, `heedb-150` pretrains over the 150-label HEEDB set + +### Environment Variables + +| Variable | Effect | +|----------|--------| +| `CACHE_DIR` | Where preprocessed datasets are cached (default `~/.cache/protossl_cache`) | +| `INVALIDATE_CACHE` | Set to `"yes im sure"` to delete the cache on import | +| `USE_HEEDB_150` | Use the 150-label HEEDB label set instead of the default (set by the `*-heedb-150-*` pretraining scripts) | +| `HIGH_MEMORY` | Load all of HEEDB into memory (~500 GB of RAM) rather than streaming it from disk (set by the HEEDB pretraining scripts) | +| `US8K_TEST_FOLD`, `IEMOCAP_TEST_FOLD` | Held-out fold for UrbanSound8K and IEMOCAP cross-validation (default `0`) | + +The ECGFounder baseline downloads its checkpoint from Hugging Face to `ecgfounder-checkpoint/` under the working directory (`scripts/` for the run scripts) on first use. + +## Contributing + +We use pre-commit hooks to maintain coding style and to enforce data stewardship. Source data should never be committed to this repo, including exploratory notebooks which may accidentally leak source data. The pre-commit hooks (which strip notebook outputs, scan for secrets, and block direct commits to `main`) do not necessarily prevent this and can be circumvented, but may help prevent obvious cases. Every tracked file must also be covered by `.github/CODEOWNERS`. + +## License + +Apache 2.0, see [`LICENSE`](LICENSE). diff --git a/configs/audio/pretrain-supervised-v2.yaml b/configs/audio/pretrain-supervised-v2.yaml new file mode 100644 index 0000000..0b91ab6 --- /dev/null +++ b/configs/audio/pretrain-supervised-v2.yaml @@ -0,0 +1,60 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoSSL-Audio + # save_dir: /path/to/runs/ + # name: exp-name + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + max_epochs: 1000 + log_every_n_steps: 1 +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 5 # check early stopping +model: + class_path: LitModel + init_args: + backbone_type: Cnn14 + conv_type: PANNS + input_channels: 1 + prototype_type: partial + partial_len: 32000 # 1-second @ 32 kHz + partial_overlap: 0.5 + prototype_h: 1 + prototype_w: 1 + n_prototypes_per_label: 12 + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + sampling_rate: 32000 + batch_size: 128 + num_workers: 12 + prefetch_factor: 2 +# pipeline_stage: learn-prototypes-supervised|project-prototypes-supervised|compute-embeddings|train-classifier diff --git a/configs/audio/pretrain-supervised.yaml b/configs/audio/pretrain-supervised.yaml index ad7c65d..caa7a95 100644 --- a/configs/audio/pretrain-supervised.yaml +++ b/configs/audio/pretrain-supervised.yaml @@ -24,7 +24,7 @@ trainer: mode: min save_last: True - class_path: PredictionWriter - max_epochs: 5 + max_epochs: 1000 log_every_n_steps: 1 optimizer: class_path: torch.optim.AdamW diff --git a/configs/audio/pretrain-unsupervised-v2.yaml b/configs/audio/pretrain-unsupervised-v2.yaml new file mode 100644 index 0000000..a91c8f1 --- /dev/null +++ b/configs/audio/pretrain-unsupervised-v2.yaml @@ -0,0 +1,63 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoSSL-Audio + # save_dir: /path/to/runs/ + # name: exp-name + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + max_epochs: 1000 + log_every_n_steps: 1 +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 5 # check early stopping +model: + class_path: LitModel + init_args: + backbone_type: Cnn14 + conv_type: PANNS + input_channels: 1 + prototype_type: partial + partial_len: 32000 # 1-second @ 32 kHz + partial_overlap: 0.5 + prototype_h: 1 + prototype_w: 1 + n_prototypes: 6324 # 527 audioset labels, comparing to 12 prototypes per label for supervised pretraining and to accomodate 5 PPL in highest cardinality downstream task + model_kwargs: '{"cola_loss_weight": 1, "clar_loss_weight": 0, "koleo_loss_weight": 1}' + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + sampling_rate: 32000 + batch_size: 128 + num_workers: 12 + prefetch_factor: 2 + data_wrapper_kwargs: '{"cola_view_seconds": 2}' +contrastive_pair_mode: cola +# pipeline_stage: learn-prototypes|project-prototypes|compute-embeddings|train-classifier diff --git a/configs/audio/pretrain-unsupervised.yaml b/configs/audio/pretrain-unsupervised.yaml index fc0b3d8..1dc05ef 100644 --- a/configs/audio/pretrain-unsupervised.yaml +++ b/configs/audio/pretrain-unsupervised.yaml @@ -24,7 +24,7 @@ trainer: mode: min save_last: True - class_path: PredictionWriter - max_epochs: 5 + max_epochs: 1000 log_every_n_steps: 1 optimizer: class_path: torch.optim.AdamW @@ -58,6 +58,6 @@ data: batch_size: 128 num_workers: 12 prefetch_factor: 2 - data_kwargs: '{"cola_view_seconds": 2}' + data_wrapper_kwargs: '{"cola_view_seconds": 2}' contrastive_pair_mode: cola+clar # pipeline_stage: learn-prototypes|project-prototypes|compute-embeddings|train-classifier diff --git a/configs/pretrain-supervised-heedb-150-1050.yaml b/configs/pretrain-supervised-heedb-150-1050.yaml new file mode 100644 index 0000000..1202550 --- /dev/null +++ b/configs/pretrain-supervised-heedb-150-1050.yaml @@ -0,0 +1,55 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoSSL + # save_dir: /path/to/runs/ + # name: exp-name + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + max_epochs: 1000 + log_every_n_steps: 1 +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 3 # check early stopping +model: + class_path: LitModel + init_args: + # backbone_type: resnet18 + # conv_type: 2D + prototype_type: partial + partial_len: 100 # 1-second @ 100 Hz + partial_overlap: 0.5 + n_prototypes_per_label: 7 # 150 labels in HEEDB to be comparable to unsupervised 1050 prototypes + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + batch_size: 512 + num_workers: 4 +# pipeline_stage: learn-prototypes-supervised|project-prototypes-supervised|compute-embeddings|train-classifier diff --git a/configs/pretrain-supervised-heedb-150.yaml b/configs/pretrain-supervised-heedb-150.yaml new file mode 100644 index 0000000..5dfd418 --- /dev/null +++ b/configs/pretrain-supervised-heedb-150.yaml @@ -0,0 +1,55 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoSSL + # save_dir: /path/to/runs/ + # name: exp-name + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + max_epochs: 1000 + log_every_n_steps: 1 +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 3 # check early stopping +model: + class_path: LitModel + init_args: + # backbone_type: resnet18 + # conv_type: 2D + prototype_type: partial + partial_len: 100 # 1-second @ 100 Hz + partial_overlap: 0.5 + n_prototypes_per_label: 50 # 150 labels in HEEDB to be comparable to unsupervised 7500 prototypes + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + batch_size: 512 + num_workers: 4 +# pipeline_stage: learn-prototypes-supervised|project-prototypes-supervised|compute-embeddings|train-classifier diff --git a/configs/pretrain-unsupervised-heedb-150-1050.yaml b/configs/pretrain-unsupervised-heedb-150-1050.yaml new file mode 100644 index 0000000..d54d3be --- /dev/null +++ b/configs/pretrain-unsupervised-heedb-150-1050.yaml @@ -0,0 +1,56 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoSSL + # save_dir: /path/to/runs/ + # name: exp-name + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + # max_epochs: 100 for learn-prototypes, 1000 for train-classifier (default if you don't set it) + log_every_n_steps: 1 +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 3 # check early stopping +model: + class_path: LitModel + init_args: + # backbone_type: resnet18 + # conv_type: 2D + prototype_type: partial + partial_len: 100 # 1-second @ 100 Hz + partial_overlap: 0.5 + n_prototypes: 1050 + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + batch_size: 512 + num_workers: 4 +contrastive_pair_mode: pclr +# pipeline_stage: learn-prototypes|project-prototypes|compute-embeddings|train-classifier diff --git a/configs/pretrain-unsupervised-heedb-150.yaml b/configs/pretrain-unsupervised-heedb-150.yaml new file mode 100644 index 0000000..cc24344 --- /dev/null +++ b/configs/pretrain-unsupervised-heedb-150.yaml @@ -0,0 +1,56 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoSSL + # save_dir: /path/to/runs/ + # name: exp-name + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + # max_epochs: 100 for learn-prototypes, 1000 for train-classifier (default if you don't set it) + log_every_n_steps: 1 +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 3 # check early stopping +model: + class_path: LitModel + init_args: + # backbone_type: resnet18 + # conv_type: 2D + prototype_type: partial + partial_len: 100 # 1-second @ 100 Hz + partial_overlap: 0.5 + n_prototypes: 7500 + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + batch_size: 512 + num_workers: 4 +contrastive_pair_mode: pclr +# pipeline_stage: learn-prototypes|project-prototypes|compute-embeddings|train-classifier diff --git a/configs/pretrain-unsupervised-no-proto.yaml b/configs/pretrain-unsupervised-no-proto.yaml new file mode 100644 index 0000000..67fa57a --- /dev/null +++ b/configs/pretrain-unsupervised-no-proto.yaml @@ -0,0 +1,52 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoSSL + # save_dir: /path/to/runs/ + # name: exp-name + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + max_epochs: 100 # for train-contraster + log_every_n_steps: 1 +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 3 # check early stopping +model: + class_path: LitModel + # init_args: + # backbone_type: resnet18 + # conv_type: 2D + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + batch_size: 512 + num_workers: 4 +contrastive_pair_mode: pclr +pipeline_stage: train-contraster diff --git a/configs/multi-arm-heedb/1D-global.yaml b/large-user-study/configs/1D-global.yaml similarity index 85% rename from configs/multi-arm-heedb/1D-global.yaml rename to large-user-study/configs/1D-global.yaml index bdd024f..d0a2850 100644 --- a/configs/multi-arm-heedb/1D-global.yaml +++ b/large-user-study/configs/1D-global.yaml @@ -12,21 +12,20 @@ trainer: callbacks: - class_path: lightning.pytorch.callbacks.EarlyStopping init_args: - monitor: val_loss - mode: min + monitor: val_ap + mode: max patience: 10 # check reduce lr on plateau - class_path: lightning.pytorch.callbacks.LearningRateMonitor init_args: logging_interval: epoch - class_path: lightning.pytorch.callbacks.ModelCheckpoint init_args: - monitor: val_loss - mode: min + monitor: val_ap + mode: max save_last: True - class_path: PredictionWriter max_epochs: 1000 log_every_n_steps: 1 - val_check_interval: 0.25 # affects hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: @@ -40,10 +39,10 @@ lr_scheduler: model: class_path: LitModel init_args: - resnet_type: resnet18 + backbone_type: resnet18 conv_type: 1D prototype_type: global - # n_prototypes_per_label: tune this! + n_prototypes_per_label: 50 # pretrained_weights: /path/to/weights data: class_path: LitData @@ -63,4 +62,4 @@ data: - VENTRICULAR TACHYCARDIA - VENTRICULAR-PACED RHYTHM - WITH SINUS ARRHYTHMIA -pipeline_stage: learn-prototypes-supervised +# pipeline_stage: diff --git a/configs/multi-arm-heedb/2D-global.yaml b/large-user-study/configs/2D-global.yaml similarity index 83% rename from configs/multi-arm-heedb/2D-global.yaml rename to large-user-study/configs/2D-global.yaml index a15df8d..0a31a36 100644 --- a/configs/multi-arm-heedb/2D-global.yaml +++ b/large-user-study/configs/2D-global.yaml @@ -12,21 +12,20 @@ trainer: callbacks: - class_path: lightning.pytorch.callbacks.EarlyStopping init_args: - monitor: val_loss - mode: min + monitor: val_ap + mode: max patience: 10 # check reduce lr on plateau - class_path: lightning.pytorch.callbacks.LearningRateMonitor init_args: logging_interval: epoch - class_path: lightning.pytorch.callbacks.ModelCheckpoint init_args: - monitor: val_loss - mode: min + monitor: val_ap + mode: max save_last: True - class_path: PredictionWriter max_epochs: 1000 log_every_n_steps: 1 - val_check_interval: 0.25 # affects hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: @@ -40,10 +39,10 @@ lr_scheduler: model: class_path: LitModel init_args: - resnet_type: resnet18 + backbone_type: resnet18 conv_type: 2D prototype_type: global - # n_prototypes_per_label: tune this! + n_prototypes_per_label: 50 # pretrained_weights: /path/to/weights data: class_path: LitData @@ -54,4 +53,4 @@ data: prefetch_factor: 4 label_subset: - NORMAL SINUS RHYTHM -pipeline_stage: learn-prototypes-supervised +# pipeline_stage: diff --git a/configs/multi-arm-heedb/2D-partial.yaml b/large-user-study/configs/2D-partial.yaml similarity index 86% rename from configs/multi-arm-heedb/2D-partial.yaml rename to large-user-study/configs/2D-partial.yaml index a9dacd1..84ff484 100644 --- a/configs/multi-arm-heedb/2D-partial.yaml +++ b/large-user-study/configs/2D-partial.yaml @@ -12,21 +12,20 @@ trainer: callbacks: - class_path: lightning.pytorch.callbacks.EarlyStopping init_args: - monitor: val_loss - mode: min + monitor: val_ap + mode: max patience: 10 # check reduce lr on plateau - class_path: lightning.pytorch.callbacks.LearningRateMonitor init_args: logging_interval: epoch - class_path: lightning.pytorch.callbacks.ModelCheckpoint init_args: - monitor: val_loss - mode: min + monitor: val_ap + mode: max save_last: True - class_path: PredictionWriter max_epochs: 1000 log_every_n_steps: 1 - val_check_interval: 0.25 # affects hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: @@ -40,12 +39,12 @@ lr_scheduler: model: class_path: LitModel init_args: - resnet_type: resnet18 + backbone_type: resnet18 conv_type: 2D prototype_type: partial partial_len: 100 # 1-second @ 100 Hz partial_overlap: 0.5 - # n_prototypes_per_label: tune this! + n_prototypes_per_label: 50 # pretrained_weights: /path/to/weights data: class_path: LitData @@ -64,4 +63,4 @@ data: - PREMATURE VENTRICULAR COMPLEXES - RIGHT AXIS DEVIATION - RIGHT BUNDLE BRANCH BLOCK -pipeline_stage: learn-prototypes-supervised +# pipeline_stage: diff --git a/large-user-study/configs/fusion.yaml b/large-user-study/configs/fusion.yaml new file mode 100644 index 0000000..6d99c97 --- /dev/null +++ b/large-user-study/configs/fusion.yaml @@ -0,0 +1,52 @@ +# lightning.pytorch==2.6.0 +seed_everything: 42 +trainer: + devices: 1 + precision: 32 + logger: + class_path: protossl.lightning_utils.StrictWandbLogger + init_args: + project: ProtoECGNet-HEEDB + # save_dir: /path/to/runs/ + name: ProtoECGNet-HEEDB-Fusion + callbacks: + - class_path: lightning.pytorch.callbacks.EarlyStopping + init_args: + monitor: val_loss + mode: min + patience: 10 # check reduce lr on plateau + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: epoch + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: val_loss + mode: min + save_last: True + - class_path: PredictionWriter + max_epochs: 1000 + log_every_n_steps: 1 + # val_check_interval: 0.25 # set for hyberband pruning +optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.001 + weight_decay: 0.01 +lr_scheduler: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + monitor: val_loss + patience: 3 # check early stopping +model: + class_path: LitModel + # init_args: + # blocks: + # pretrained_weights: /path/to/weights +data: + class_path: LitData + init_args: + # dataset_path: /path/to/dataset + batch_size: 512 + num_workers: 4 + prefetch_factor: 4 +# pipeline_stage: diff --git a/large-user-study/configs/lambda-sweep/wandb-sweep.yaml b/large-user-study/configs/lambda-sweep/wandb-sweep.yaml new file mode 100644 index 0000000..b69c43d --- /dev/null +++ b/large-user-study/configs/lambda-sweep/wandb-sweep.yaml @@ -0,0 +1,56 @@ +program: protossl.protoecgnet_trainer +# name: ProtoECGNet-HEEDB-1D-Global-Coef +# name: ProtoECGNet-HEEDB-2D-Global-Coef +name: ProtoECGNet-HEEDB-2D-Partial-Coef +method: bayes +metric: + goal: maximize + name: val_ap +parameters: + model.lam_clst: + distribution: log_uniform_values + min: 0.001 + max: 1.0 + model.lam_sep: + distribution: log_uniform_values + min: 0.0001 + max: 0.2 + model.lam_div: + distribution: log_uniform_values + min: 10.0 + max: 1000.0 + model.lam_cntrst: + distribution: log_uniform_values + min: 1.0 + max: 3000.0 +early_terminate: + # val_ap is configured to log/step at the end of every train epoch + # this hyperband config gives the step brackets [8, 4, 2] + # so the prune happens at 2, 4, 8 train epochs + type: hyperband + s: 3 + eta: 2 + max_iter: 16 +command: + - ${env} + - HIGH_MEMORY=1 + - python + - "-m" + - ${program} + - "--pipeline-stage" + - "learn-prototypes-supervised" + - "--config" + # - "../1D-global.yaml" + # - "../2D-global.yaml" + - "../2D-partial.yaml" + - "--trainer.logger.save_dir" + - "/home/songs1/outputs-tune-heedb" + - "--data.dataset_path" + - "/opt/gpudata/ecg/heedb" + - "--data.data_kwargs" + - "{\"heedb_split_type\": \"by-label\"}" + - "--data.num_workers" + - "8" + - "--data.prefetch_factor" + - "4" + - ${args} diff --git a/configs/multi-arm-heedb/wandb-sweep.yaml b/large-user-study/configs/ppl-sweep/wandb-sweep.yaml similarity index 100% rename from configs/multi-arm-heedb/wandb-sweep.yaml rename to large-user-study/configs/ppl-sweep/wandb-sweep.yaml diff --git a/large-user-study/notebooks/confirm-new-heedb-splits.ipynb b/large-user-study/notebooks/confirm-new-heedb-splits.ipynb new file mode 100644 index 0000000..0669ef8 --- /dev/null +++ b/large-user-study/notebooks/confirm-new-heedb-splits.ipynb @@ -0,0 +1,227 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "849ac751", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import torch\n", + "\n", + "from protossl.datasets import HeedbECGDataset\n", + "\n", + "dataset_path = \"/opt/gpudata/ecg/heedb\"" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "83c72a24", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "================get_heedb_metadata=================\n", + "reading HEEDB metadata\n", + "read MGB data\n", + "read Emory data\n", + "===================================================\n", + "=================make_heedb_labels=================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:05<00:00, 1781953.72it/s]\n", + "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 3226768.39it/s]\n", + "Converting code string to labels: 100%|██████████| 8163230/8163230 [01:17<00:00, 105109.73it/s]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 8163230 ECGs and 11607259 annotations, 8163229 matched (1 ECG annotations were missing and filled with 0s)\n", + "===================================================\n", + "===============StreamingECGWaveforms===============\n", + "Using streaming ECG waveforms, will load and transform data on the fly\n", + "===================================================\n", + "=================make_heedb_labels=================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Converting code string to labels: 100%|██████████| 453512/453512 [00:02<00:00, 225403.89it/s]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 453512 ECGs and 11607259 annotations, 453512 matched (0 ECG annotations were missing and filled with 0s)\n", + "===================================================\n", + "===============StreamingECGWaveforms===============\n", + "Using streaming ECG waveforms, will load and transform data on the fly\n", + "===================================================\n", + "=================make_heedb_labels=================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Converting code string to labels: 100%|██████████| 453512/453512 [00:04<00:00, 108876.87it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 453512 ECGs and 11607259 annotations, 453512 matched (0 ECG annotations were missing and filled with 0s)\n", + "===================================================\n", + "===============StreamingECGWaveforms===============\n", + "Using streaming ECG waveforms, will load and transform data on the fly\n", + "===================================================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "ds_train = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"train\",\n", + " sampling_rate=100,\n", + " heedb_split_type=\"by-label\",\n", + ")\n", + "ds_val = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"val\",\n", + " sampling_rate=100,\n", + " heedb_split_type=\"by-label\",\n", + ")\n", + "ds_test = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"test\",\n", + " sampling_rate=100,\n", + " heedb_split_type=\"by-label\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e7b62f26", + "metadata": {}, + "outputs": [], + "source": [ + "def find_tol(pairs: list[tuple[torch.Tensor, torch.Tensor]]) -> float:\n", + " last_tol = np.nan\n", + " tols = np.outer(10.0 ** -np.arange(1, 5), (1, 0.5, 0.25)).ravel() # [0.1, 0.05, 0.025, 0.01, ...]\n", + " for tol in tols:\n", + " try:\n", + " for a, b in pairs:\n", + " assert np.isclose(a, b, atol=tol).all()\n", + " except AssertionError:\n", + " return last_tol\n", + " last_tol = tol\n", + " return last_tol" + ] + }, + { + "cell_type": "code", + "execution_count": 42, + "id": "c44f839c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Label prevalences are preserved to 0.1% tolerance\n" + ] + } + ], + "source": [ + "train_prev = ds_train.labels.sum(axis=0) / len(ds_train)\n", + "val_prev = ds_val.labels.sum(axis=0) / len(ds_val)\n", + "test_prev = ds_test.labels.sum(axis=0) / len(ds_test)\n", + "\n", + "tol = find_tol([\n", + " (train_prev, val_prev),\n", + " (train_prev, test_prev),\n", + " (val_prev, test_prev),\n", + "])\n", + "print(f\"Label prevalences are preserved to {tol*100:0.1f}% tolerance\")" + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "id": "903e7305", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Cooccurrence frequencies are preserved to 1.0% tolerance\n" + ] + } + ], + "source": [ + "train_cooc = ds_train.get_label_cooccurrence()\n", + "val_cooc = ds_val.get_label_cooccurrence()\n", + "test_cooc = ds_test.get_label_cooccurrence()\n", + "\n", + "tol = find_tol([\n", + " (train_cooc, val_cooc),\n", + " (train_cooc, test_cooc),\n", + " (val_cooc, test_cooc),\n", + "])\n", + "print(f\"Cooccurrence frequencies are preserved to {tol*100:0.1f}% tolerance\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "150bd5e4", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "protossl", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/large-user-study/notebooks/muse-accuracy.ipynb b/large-user-study/notebooks/muse-accuracy.ipynb new file mode 100644 index 0000000..7f2ee89 --- /dev/null +++ b/large-user-study/notebooks/muse-accuracy.ipynb @@ -0,0 +1,585 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "308c1648", + "metadata": {}, + "outputs": [], + "source": [ + "from protossl.datasets import HeedbECGDataset\n", + "from protossl.defines import HEEDB_TARGETS\n", + "from omegaconf import OmegaConf\n", + "from pathlib import Path\n", + "import numpy as np\n", + "import pandas as pd\n", + "from sklearn.metrics import roc_curve, balanced_accuracy_score, roc_auc_score\n", + "\n", + "dataset_path = \"/opt/gpudata/ecg/heedb\"\n", + "arms = [\"2D-partial\", \"2D-global\", \"1D-global\"]\n", + "run_path = Path(\"/opt/gpu_working/steven/protoecgnet-heedb\")\n", + "config_path = Path(\"/opt/gpu_working/steven/ProtoSSL/large-user-study/configs\")\n", + "# dropping a branch from the head drops the labels only it covered, so the probs\n", + "# are scored against labels.yaml rather than all of HEEDB_TARGETS\n", + "fusion_path = run_path / \"fusion-tuned/tune-fusion-classifier-drop-2D-global\"\n", + "fusion_labels = list(\n", + " OmegaConf.load(fusion_path / \"labels.yaml\").data.init_args.label_subset\n", + ")\n", + "# fusion label order is a subsequence of HEEDB_TARGETS, i.e. of the dataset columns\n", + "label_idxs = np.asarray([list(HEEDB_TARGETS).index(l) for l in fusion_labels])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9edcebae", + "metadata": {}, + "outputs": [], + "source": [ + "def stack_branched_arms():\n", + " arm_labels = [l for arm in arms for l in OmegaConf.load(config_path / f\"{arm}.yaml\").data.init_args.label_subset]\n", + " idxs = np.asarray([arm_labels.index(l) for l in fusion_labels]) # indices to normalize to fusion order\n", + "\n", + " ret = dict()\n", + " for split in [\"val\", \"test\"] :\n", + " arm_probs = {arm: np.load(run_path / arm / f\"train-classifier/latest/{split}_probs.npy\") for arm in arms}\n", + " stacked_probs = np.concat([arm_probs[arm] for arm in arms], axis=1)\n", + " ret[split] = stacked_probs[:, idxs]\n", + " return ret" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "68d44cd8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "================get_heedb_metadata=================\n", + "using by-label splits\n", + "reading HEEDB metadata from on-disk cache: /home/songs1/.cache/protossl_cache/abac20ad.csv\n", + "===================================================\n", + "=================make_heedb_labels=================\n", + "reading HEEDB labels from on-disk cache: /home/songs1/.cache/protossl_cache/51ff3555.npy\n", + "replaying below stats from cache:\n", + "Of 453512 ECGs and 11607259 annotations, 453512 matched (0 ECG annotations were missing and filled with 0s)\n", + "===================================================\n", + "===============StreamingECGWaveforms===============\n", + "Using streaming ECG waveforms, will load and transform data on the fly\n", + "===================================================\n", + "=================make_heedb_labels=================\n", + "reading HEEDB labels from on-disk cache: /home/songs1/.cache/protossl_cache/c6af0872.npy\n", + "replaying below stats from cache:\n", + "Of 453512 ECGs and 11607259 annotations, 453512 matched (0 ECG annotations were missing and filled with 0s)\n", + "===================================================\n", + "===============StreamingECGWaveforms===============\n", + "Using streaming ECG waveforms, will load and transform data on the fly\n", + "===================================================\n", + "=================make_heedb_labels=================\n", + "reading HEEDB labels from on-disk cache: /home/songs1/.cache/protossl_cache/0a73c965.npy\n", + "replaying below stats from cache:\n", + "Of 453512 ECGs and 11607259 annotations, 453512 matched (0 ECG annotations were missing and filled with 0s)\n", + "===================================================\n", + "===============StreamingECGWaveforms===============\n", + "Using streaming ECG waveforms, will load and transform data on the fly\n", + "===================================================\n" + ] + } + ], + "source": [ + "ds_phys = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"test\",\n", + " sampling_rate=100,\n", + " label_src=\"original_physician\",\n", + " heedb_split_type=\"by-label\",\n", + ")\n", + "\n", + "ds_phys_val = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"val\",\n", + " sampling_rate=100,\n", + " label_src=\"original_physician\",\n", + " heedb_split_type=\"by-label\",\n", + ")\n", + "\n", + "ds_muse = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"test\",\n", + " sampling_rate=100,\n", + " label_src=\"original_muse\",\n", + " heedb_split_type=\"by-label\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "55a31c9e", + "metadata": {}, + "outputs": [], + "source": [ + "y_true = ds_phys.labels.numpy()[:, label_idxs]\n", + "y_true_val = ds_phys_val.labels.numpy()[:, label_idxs]\n", + "y_pred_muse = ds_muse.labels.numpy()[:, label_idxs]" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a84dca44", + "metadata": {}, + "outputs": [], + "source": [ + "def compute_threshold(_y_true, _y_prob):\n", + " # NOTE: should be computed on either train/val set, NOT test\n", + " fpr, tpr, thresholds = roc_curve(_y_true, _y_prob)\n", + " j_scores = tpr - fpr\n", + " optimal_idx = np.argmax(j_scores)\n", + " optimal_threshold = thresholds[optimal_idx]\n", + " return optimal_threshold" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "be9ee773", + "metadata": {}, + "outputs": [], + "source": [ + "# temp = stack_branched_arms()\n", + "# y_prob_branched, y_prob_branched_val = temp[\"test\"], temp[\"val\"]\n", + "y_prob_fusion = np.load(fusion_path / \"test_probs.npy\")\n", + "y_prob_fusion_val = np.load(fusion_path / \"val_probs.npy\")\n", + "assert y_prob_fusion.shape[1] == len(fusion_labels)\n", + "\n", + "# y_pred_branched = np.zeros_like(y_pred_muse)\n", + "y_pred_fusion = np.zeros_like(y_pred_muse)\n", + "accuracy_results = []\n", + "roc_results = []\n", + "for i, label in enumerate(fusion_labels):\n", + " # branched_threshold = compute_threshold(y_true_val[:, i], y_prob_branched_val[:, i])\n", + " # y_pred_branched[:, i] = (y_prob_branched[:, i] > branched_threshold).astype(int)\n", + " fusion_threshold = compute_threshold(y_true_val[:, i], y_prob_fusion_val[:, i])\n", + " y_pred_fusion[:, i] = (y_prob_fusion[:, i] > fusion_threshold).astype(int)\n", + " label_acc_results = {\n", + " \"label\": label,\n", + " \"muse: BalAcc\": balanced_accuracy_score(y_true[:, i], y_pred_muse[:, i]),\n", + " # \"branched: BalAcc\": balanced_accuracy_score(y_true[:, i], y_pred_branched[:, i]),\n", + " \"fusion: BalAcc\": balanced_accuracy_score(y_true[:, i], y_pred_fusion[:, i]),\n", + " }\n", + " label_roc_results = {\n", + " \"label\": label,\n", + " # NOTE: can't compute muse roc given no probs\n", + " # \"branched: AUROC\": roc_auc_score(y_true[:, i], y_prob_branched[:, i]),\n", + " \"fusion: AUROC\": roc_auc_score(y_true[:, i], y_prob_fusion[:, i]),\n", + " }\n", + " accuracy_results.append(label_acc_results)\n", + " roc_results.append(label_roc_results)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "23b6c024", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
labelfusion: AUROC
0ANTERIOR INFARCT0.953
1ATRIAL FIBRILLATION0.959
2ATRIAL FLUTTER0.953
3ATRIAL-PACED RHYTHM0.986
4INCOMPLETE RIGHT BUNDLE BRANCH BLOCK0.962
5INFERIOR INFARCT0.976
6LATERAL INFARCT0.957
7LEFT BUNDLE BRANCH BLOCK0.987
8PREMATURE ATRIAL COMPLEXES0.912
9PREMATURE VENTRICULAR COMPLEXES0.896
10RIGHT AXIS DEVIATION0.982
11RIGHT BUNDLE BRANCH BLOCK0.986
12SINUS BRADYCARDIA0.981
13SINUS RHYTHM0.833
14SINUS TACHYCARDIA0.981
15VENTRICULAR TACHYCARDIA0.856
16VENTRICULAR-PACED RHYTHM0.982
17WITH 1ST DEGREE AV BLOCK0.978
18WITH SINUS ARRHYTHMIA0.973
19Macro Average0.952
\n", + "
" + ], + "text/plain": [ + " label fusion: AUROC\n", + "0 ANTERIOR INFARCT 0.953\n", + "1 ATRIAL FIBRILLATION 0.959\n", + "2 ATRIAL FLUTTER 0.953\n", + "3 ATRIAL-PACED RHYTHM 0.986\n", + "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.962\n", + "5 INFERIOR INFARCT 0.976\n", + "6 LATERAL INFARCT 0.957\n", + "7 LEFT BUNDLE BRANCH BLOCK 0.987\n", + "8 PREMATURE ATRIAL COMPLEXES 0.912\n", + "9 PREMATURE VENTRICULAR COMPLEXES 0.896\n", + "10 RIGHT AXIS DEVIATION 0.982\n", + "11 RIGHT BUNDLE BRANCH BLOCK 0.986\n", + "12 SINUS BRADYCARDIA 0.981\n", + "13 SINUS RHYTHM 0.833\n", + "14 SINUS TACHYCARDIA 0.981\n", + "15 VENTRICULAR TACHYCARDIA 0.856\n", + "16 VENTRICULAR-PACED RHYTHM 0.982\n", + "17 WITH 1ST DEGREE AV BLOCK 0.978\n", + "18 WITH SINUS ARRHYTHMIA 0.973\n", + "19 Macro Average 0.952" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "roc_df = pd.DataFrame(roc_results)\n", + "idx = len(roc_df)\n", + "roc_df.loc[idx, \"label\"] = \"Macro Average\"\n", + "for c in [c for c in roc_df.columns if c != \"label\"]:\n", + " roc_df.loc[idx, c] = roc_df[c].mean()\n", + "\n", + "with pd.option_context(\"display.precision\", 3):\n", + " display(roc_df)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "c2c33568", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
labelmuse: BalAccfusion: BalAcc
0ANTERIOR INFARCT0.9490.893
1ATRIAL FIBRILLATION0.9210.905
2ATRIAL FLUTTER0.8580.888
3ATRIAL-PACED RHYTHM0.8730.943
4INCOMPLETE RIGHT BUNDLE BRANCH BLOCK0.9100.908
5INFERIOR INFARCT0.9720.918
6LATERAL INFARCT0.9730.909
7LEFT BUNDLE BRANCH BLOCK0.9400.958
8PREMATURE ATRIAL COMPLEXES0.9100.850
9PREMATURE VENTRICULAR COMPLEXES0.9290.831
10RIGHT AXIS DEVIATION0.9100.956
11RIGHT BUNDLE BRANCH BLOCK0.9690.962
12SINUS BRADYCARDIA0.9570.956
13SINUS RHYTHM0.8510.757
14SINUS TACHYCARDIA0.9580.954
15VENTRICULAR TACHYCARDIA0.6490.773
16VENTRICULAR-PACED RHYTHM0.7380.948
17WITH 1ST DEGREE AV BLOCK0.9670.938
18WITH SINUS ARRHYTHMIA0.9570.931
19Macro Average0.9050.904
\n", + "
" + ], + "text/plain": [ + " label muse: BalAcc fusion: BalAcc\n", + "0 ANTERIOR INFARCT 0.949 0.893\n", + "1 ATRIAL FIBRILLATION 0.921 0.905\n", + "2 ATRIAL FLUTTER 0.858 0.888\n", + "3 ATRIAL-PACED RHYTHM 0.873 0.943\n", + "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.910 0.908\n", + "5 INFERIOR INFARCT 0.972 0.918\n", + "6 LATERAL INFARCT 0.973 0.909\n", + "7 LEFT BUNDLE BRANCH BLOCK 0.940 0.958\n", + "8 PREMATURE ATRIAL COMPLEXES 0.910 0.850\n", + "9 PREMATURE VENTRICULAR COMPLEXES 0.929 0.831\n", + "10 RIGHT AXIS DEVIATION 0.910 0.956\n", + "11 RIGHT BUNDLE BRANCH BLOCK 0.969 0.962\n", + "12 SINUS BRADYCARDIA 0.957 0.956\n", + "13 SINUS RHYTHM 0.851 0.757\n", + "14 SINUS TACHYCARDIA 0.958 0.954\n", + "15 VENTRICULAR TACHYCARDIA 0.649 0.773\n", + "16 VENTRICULAR-PACED RHYTHM 0.738 0.948\n", + "17 WITH 1ST DEGREE AV BLOCK 0.967 0.938\n", + "18 WITH SINUS ARRHYTHMIA 0.957 0.931\n", + "19 Macro Average 0.905 0.904" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "acc_df = pd.DataFrame(accuracy_results)\n", + "idx = len(acc_df)\n", + "acc_df.loc[idx, \"label\"] = \"Macro Average\"\n", + "for c in [c for c in acc_df.columns if c != \"label\"]:\n", + " acc_df.loc[idx, c] = acc_df[c].mean()\n", + "\n", + "with pd.option_context(\"display.precision\", 3):\n", + " display(acc_df)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa320886", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "protossl", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/large-user-study/notebooks/spoof-heedb.ipynb b/large-user-study/notebooks/spoof-heedb.ipynb new file mode 100644 index 0000000..f3d84bb --- /dev/null +++ b/large-user-study/notebooks/spoof-heedb.ipynb @@ -0,0 +1,229 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "5a3c4ae6", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "from pathlib import Path\n", + "\n", + "in_path = Path(\"/opt/gpudata/ecg/heedb\")\n", + "out_path = Path(\"/opt/gpudata/ecg/heedb-spoof\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "46892b19", + "metadata": {}, + "outputs": [], + "source": [ + "mgh_meta = pd.read_csv(in_path / \"I0001/metadata/metadata.csv\", usecols=[\"BDSPPatientID\", \"SexDSC\", \"AgeAtAcquisition\", \"FileName\"])\n", + "euh_meta = pd.read_csv(in_path / \"I0006/metadata/metadata.csv\", usecols=[\"BDSPPatientID\", \"Sex\", \"AgeAtAcquisition\", \"FileName\"])\n", + "mgh_labels = pd.read_csv(in_path / \"I0001/12SL_diagnoses/diagnoses_acquisition.csv\")\n", + "euh_labels = pd.read_csv(in_path / \"I0006/12SL_diagnoses/diagnoses_acquisition.csv\")" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "57f998ca", + "metadata": {}, + "outputs": [], + "source": [ + "mgh_years = mgh_meta[\"FileName\"].str.split(\"/\").str[2]\n", + "euh_years = euh_meta[\"FileName\"].str.split(\"/\").str[1]" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "f634cd61", + "metadata": {}, + "outputs": [], + "source": [ + "mgh_years = mgh_years.astype(int)\n", + "euh_years = euh_years.astype(int)" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "id": "381e5254", + "metadata": {}, + "outputs": [], + "source": [ + "mgh_slim = mgh_meta[mgh_years.isin({2010, 2021, 2022})]\n", + "euh_slim = euh_meta[euh_years == 2010]" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "id": "70c3e6ba", + "metadata": {}, + "outputs": [], + "source": [ + "mgh_slim_set = set(mgh_slim[\"FileName\"])\n", + "euh_slim_set = set(euh_slim[\"FileName\"])" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "id": "ea368935", + "metadata": {}, + "outputs": [], + "source": [ + "mgh_slim_labels = mgh_labels[mgh_labels[\"FileName\"].isin(mgh_slim_set)]\n", + "euh_slim_labels = euh_labels[euh_labels[\"FileName\"].isin(euh_slim_set)]" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "id": "172d912a", + "metadata": {}, + "outputs": [], + "source": [ + "mgh_slim.to_csv(out_path / \"I0001/metadata/metadata.csv\", index=False)\n", + "euh_slim.to_csv(out_path / \"I0006/metadata/metadata.csv\", index=False)\n", + "mgh_slim_labels.to_csv(out_path / \"I0001/12SL_diagnoses/diagnoses_acquisition.csv\", index=False)\n", + "euh_slim_labels.to_csv(out_path / \"I0006/12SL_diagnoses/diagnoses_acquisition.csv\", index=False)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "b1b3af76", + "metadata": {}, + "outputs": [], + "source": [ + "from protossl.datasets import HeedbECGDataset" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "93e7d52a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "================get_heedb_metadata=================\n", + "reading HEEDB metadata\n", + "read MGB data\n", + "read Emory data\n", + "===================================================\n", + "=================make_heedb_labels=================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Creating mgb fname to code mapping: 100%|██████████| 1018374/1018374 [00:00<00:00, 2941186.29it/s]\n", + "Creating emory fname to code mapping: 100%|██████████| 27371/27371 [00:00<00:00, 3601866.62it/s]\n", + "Converting code string to labels: 100%|██████████| 59002/59002 [00:00<00:00, 537972.79it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 59002 ECGs and 1045745 annotations, 59002 matched (0 ECG annotations were missing and filled with 0s)\n", + "===================================================\n", + "===============StreamingECGWaveforms===============\n", + "Using streaming ECG waveforms, will load and transform data on the fly\n", + "===================================================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "ds = HeedbECGDataset(\n", + " dataset_path=\"/opt/gpudata/ecg/heedb-spoof\",\n", + " split=\"train\",\n", + " sampling_rate=100,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "d0aafcdb", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "62714a78", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([[0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " ...,\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.],\n", + " [0., 0., 0., ..., 0., 0., 0.]], shape=(5120, 20), dtype=float32)" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "np.load(\"temp/fusion/train-fusion-classifier/latest/probs.npy\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a7bd307c", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "protossl", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/large-user-study/scripts/_cache_heedb_overread.py b/large-user-study/scripts/_cache_heedb_overread.py new file mode 100644 index 0000000..b666dc2 --- /dev/null +++ b/large-user-study/scripts/_cache_heedb_overread.py @@ -0,0 +1,18 @@ +import os + +os.environ["HIGH_MEMORY"] = "1" + +from protossl.datasets import HeedbECGDataset + +DATASET_PATH = "/opt/gpudata/ecg/heedb" + +# builds the overread mask and waveform caches up front, so the concurrently +# launched arm jobs only ever read them +for split in ["train", "val", "test"]: + HeedbECGDataset( + dataset_path=DATASET_PATH, + sampling_rate=100, + split=split, + heedb_split_type="by-label", + drop_blank_overread=True, + ) diff --git a/large-user-study/scripts/_eval_probs_bootstrapped.py b/large-user-study/scripts/_eval_probs_bootstrapped.py new file mode 100644 index 0000000..0fc64be --- /dev/null +++ b/large-user-study/scripts/_eval_probs_bootstrapped.py @@ -0,0 +1,161 @@ +import os +from argparse import ArgumentParser +from ast import literal_eval +from collections import defaultdict + +import numpy as np +import pandas as pd +import scipy.stats as st +from omegaconf import OmegaConf +from pqdm.processes import pqdm +from sklearn.metrics import average_precision_score, roc_auc_score + +from protossl.datasets import EchoNextECGDataset, infer_dataset_class_from_path +from protossl.defines import ECHONEXT_COMPOSITE_TARGET + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--dataset-path", required=True) + parser.add_argument("--data-kwargs", required=True) + parser.add_argument("--probs-npy", required=True) + parser.add_argument("--output-path", required=True) + group = parser.add_mutually_exclusive_group() + group.add_argument("--label-subset", nargs="+") + group.add_argument("--label-subset-config") + parser.add_argument("--n-bootstraps", type=int, default=1000) + parser.add_argument("--n-jobs", type=int, default=24) + args = parser.parse_args() + args.data_kwargs = literal_eval(args.data_kwargs) + return args + + +def worker( + y_test: np.ndarray, + y_prob: np.ndarray, + n_bootstraps: int, +) -> dict: + bootstrap_n = len(y_test) + label_pos_frac = y_test.sum() / len(y_test) + label_metrics = dict() + pos_idxs = np.argwhere(y_test).squeeze(1) + neg_idxs = np.argwhere(1 - y_test).squeeze(1) + + bootstrapped_metrics = defaultdict(list) + for b in range(n_bootstraps): + rng = np.random.default_rng(b) + bootstrap_n_pos = int(bootstrap_n * label_pos_frac) + bootstrap_n_pos = max(bootstrap_n_pos, 1) # require at least 1 sample + bootstrap_n_neg = bootstrap_n - bootstrap_n_pos + + bootstrap_pos_idxs = rng.choice(pos_idxs, bootstrap_n_pos, replace=True) + bootstrap_neg_idxs = rng.choice(neg_idxs, bootstrap_n_neg, replace=True) + + bootstrap_idxs = np.concat([bootstrap_pos_idxs, bootstrap_neg_idxs]) + + bootstrap_y_test = y_test[bootstrap_idxs] + bootstrap_y_prob = y_prob[bootstrap_idxs] + + assert bootstrap_y_test.sum() > 0 + + bootstrapped_metrics["AUROC"].append( + roc_auc_score(bootstrap_y_test, bootstrap_y_prob) + ) + bootstrapped_metrics["AUPRC"].append( + average_precision_score(bootstrap_y_test, bootstrap_y_prob) + ) + + for k, vs in bootstrapped_metrics.items(): + a = np.asarray(vs) + a_avg = np.mean(a) + lo, hi = st.t.interval(0.95, len(a) - 1, loc=a_avg, scale=st.sem(a)) + label_metrics[k] = a_avg + label_metrics[f"{k} 95% CI (lo)"] = lo + label_metrics[f"{k} 95% CI (hi)"] = hi + + return label_metrics + + +def main( + *, # enforce kwargs + dataset_path: str, + data_kwargs: dict, + probs_npy: str, + output_path: str, + label_subset: list[str] | None = None, + n_bootstraps: int = 1000, + n_jobs: int = 24, +): + ds_cls, src_label_names, is_audio = infer_dataset_class_from_path(dataset_path) + if is_audio: + raise ValueError("This eval script is meant for ECG related work") + test_ds = ds_cls( + dataset_path=dataset_path, + split="test", + sampling_rate=100, + label_subset=label_subset, + **data_kwargs, + ) + assert test_ds.labels is not None and src_label_names is not None + if label_subset is not None: + label_names = label_subset + missing = set(label_subset) - set(src_label_names) + if len(missing) > 0: + raise ValueError( + f"{len(missing)} labels in subset not in source labels: {missing}" + ) + else: + label_names = src_label_names + + composite_target = None + if ds_cls == EchoNextECGDataset: + composite_target = ECHONEXT_COMPOSITE_TARGET + + test_targets = test_ds.labels.numpy() + target_probs = np.load(probs_npy, allow_pickle=True) + assert ( + test_targets.shape == target_probs.shape + ), f"Mismatched target/predicted shapes: {test_targets.shape} vs {target_probs.shape}" + + kwargs = [ + { + # the test targets in this version are already aligned to the label subset order + "y_test": test_targets[:, i], + "y_prob": target_probs[:, i], + "n_bootstraps": n_bootstraps, + } + for i in range(len(label_names)) + ] + + results = pqdm(kwargs, worker, argument_type="kwargs", n_jobs=n_jobs) # type: ignore + metrics = {label_name: results[l] for l, label_name in enumerate(label_names)} + + metrics = pd.DataFrame.from_dict(metrics, orient="index") + metrics.index.name = "Label" + + _labels = [x for x in label_names if x != composite_target] + metrics.loc["Multilabel Averaged"] = metrics.loc[_labels].mean() + + metrics_path = os.path.join(output_path, "metrics-bootstrapped-v2.csv") + metrics.to_csv(metrics_path) + print(f"Saved metrics to {metrics_path}") + + +if __name__ == "__main__": + args = parse_args() + + # NOTE: overwrites label_subset with that found in config + if args.label_subset_config is not None: + c = OmegaConf.load(args.label_subset_config) + OmegaConf.set_struct(c, True) + args.label_subset = c.data.init_args.label_subset + + main( + dataset_path=args.dataset_path, + data_kwargs=args.data_kwargs, + probs_npy=args.probs_npy, + output_path=args.output_path, + label_subset=args.label_subset, + n_bootstraps=args.n_bootstraps, + n_jobs=args.n_jobs, + ) diff --git a/large-user-study/scripts/_tune_fusion_classifier.py b/large-user-study/scripts/_tune_fusion_classifier.py new file mode 100644 index 0000000..434d4cc --- /dev/null +++ b/large-user-study/scripts/_tune_fusion_classifier.py @@ -0,0 +1,489 @@ +""" +Tune the ProtoECGNetFusion classification head (stage 3) over precomputed +prototype similarities from `--pipeline-stage compute-fusion-embeddings`. + +The encoders are frozen in train-fusion-classifier, so the head only ever sees +the prototype similarities. Training directly on the cached similarities is +numerically the same objective as ProtoECGNetFusion.forward but skips the +dataloaders and encoder forward passes, so an Optuna study is cheap to run. + +Tuned: lam_l1 (masked L1 on off-class weights), lr, weight_decay, batch_size, +pos_weight_exp (label weights are raised to this power, 0 = unweighted, +1 = the prevalence inversion the Lightning stages use). + +--drop-branches ablates whole branches out of the head without recomputing the +embeddings: the cached similarities are laid out one contiguous block per label +(ProtoECGNetFusion.reverse_mask order), so a dropped branch is a set of columns +to skip, plus the labels only that branch covered. +Fixed (matching configs/fusion.yaml): AdamW, ReduceLROnPlateau(val_loss, +factor=0.1, patience=3), bias-free linear head with masked init (1 / -0.5). + +Model selection within a trial and across trials uses val macro AP. val_loss is +not comparable across trials - it includes the lam_l1-scaled penalty and its +classification term is pos_weight_exp-weighted (it is still used for the LR +schedule within a trial, as in the Lightning config). +""" + +import json +import os +from argparse import ArgumentParser + +import numpy as np +import optuna +import torch +import torch.nn.functional as F +from omegaconf import OmegaConf +from torchmetrics.functional.classification import ( + multilabel_auroc, + multilabel_average_precision, +) +from tqdm import trange + +from protossl.datasets import HeedbECGDataset, infer_dataset_class_from_path +from protossl.defines import SPLIT_T +from protossl.models._protoecgnet import BranchCfg, fusion_prototype_assignment + +torch.set_float32_matmul_precision("medium") + +# baseline from configs/fusion.yaml + ProtoECGNetFusion defaults, enqueued as trial 0 +BASELINE_PARAMS = { + "lam_l1": 1e-4, + "lr": 1e-3, + "weight_decay": 0.01, + "batch_size": 512, + "pos_weight_exp": 1.0, +} + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument( + "--embeds-dir", + required=True, + help="compute-fusion-embeddings log dir (has config.yaml and {split}_embeds.npy)", + ) + parser.add_argument("--output-dir", required=True) + parser.add_argument( + "--drop-branches", + nargs="+", + default=[], + metavar="NAME", + help=( + "branch names (as given to compute-fusion-embeddings) to ablate out of " + "the head: their prototype columns are skipped when loading the cached " + "similarities and any label only covered by a dropped branch is dropped " + "from the head. Needs a fresh --output-dir (the study is not comparable)" + ), + ) + parser.add_argument("--n-trials", type=int, default=50) + parser.add_argument("--max-epochs", type=int, default=100) + parser.add_argument("--early-stopping-patience", type=int, default=10) + parser.add_argument("--lr-patience", type=int, default=3) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", default="cuda") + parser.add_argument( + "--ap-tolerance", + type=float, + default=0.0, + help=( + "select the sparsest head (lowest off-class weight fraction) among " + "trials within this val AP of the best trial; 0 = pure best val AP" + ), + ) + parser.add_argument( + "--export-only", + action="store_true", + help="skip tuning, (re)select from existing study and export probs", + ) + return parser.parse_args() + + +def load_run_config(embeds_dir: str, drop_branches: list[str]) -> dict: + c = OmegaConf.load(os.path.join(embeds_dir, "config.yaml")) + assert c.pipeline_stage == "compute-fusion-embeddings", c.pipeline_stage + data_args = c.data.init_args + dataset_path: str = data_args.dataset_path + ds_cls, label_names, _ = infer_dataset_class_from_path(dataset_path) + if ds_cls != HeedbECGDataset: + raise ValueError(f"Only HEEDB is supported, got: {ds_cls}") + assert label_names is not None + label_subset = data_args.label_subset + if label_subset is not None: + label_subset = list(label_subset) + label_names = label_subset + branches = [ + BranchCfg(**b) # type: ignore + for b in OmegaConf.to_container(c.model.init_args.branches) # type: ignore + ] + unknown = set(drop_branches) - {b.name for b in branches} + if len(unknown) > 0: + raise ValueError( + f"--drop-branches has no such branch: {sorted(unknown)} " + f"(embeds run has {[b.name for b in branches]})" + ) + # the cached similarities are in reverse_mask order: one contiguous block of + # n_prototypes_per_label columns per label, in label_names order, where the + # block size comes from the branch that label was trained in + label_to_ppl = { + label: branch.n_prototypes_per_label + for branch in branches + for label in branch.label_subset + } + dropped_labels = { + label + for branch in branches + if branch.name in drop_branches + for label in branch.label_subset + } + keep_cols, keep_labels, n_cols_cached = [], [], 0 + for label in label_names: + chunk = label_to_ppl[label] + if label not in dropped_labels: + keep_cols.extend(range(n_cols_cached, n_cols_cached + chunk)) + keep_labels.append(label) + n_cols_cached += chunk + if len(keep_labels) == 0: + raise ValueError("--drop-branches would drop every label") + if len(dropped_labels) > 0: + print( + f"dropping branches {drop_branches}: head keeps " + f"{len(keep_labels)}/{len(label_names)} labels and " + f"{len(keep_cols)}/{n_cols_cached} prototype columns, " + f"dropped labels: {sorted(dropped_labels)}" + ) + branches = [b for b in branches if b.name not in drop_branches] + # keep_labels is in cached column order, so it doubles as the dataset + # label_subset (which sets both which label columns load and their order) + label_subset, label_names = keep_labels, keep_labels + return { + "dataset_path": dataset_path, + "sampling_rate": data_args.sampling_rate, + "data_kwargs": OmegaConf.to_container(data_args.data_kwargs), + "label_subset": label_subset, + "label_names": label_names, + "n_cols_cached": n_cols_cached, + # None = load every cached column + "keep_cols": np.asarray(keep_cols) if len(dropped_labels) > 0 else None, + "assign": fusion_prototype_assignment(label_names, branches), # (L, R) + } + + +def load_dataset(run_cfg: dict, split: SPLIT_T) -> HeedbECGDataset: + # with HIGH_MEMORY unset (checked in main), waveforms are streamed lazily so + # only labels/metadata are materialized - we only need the labels/weights + return HeedbECGDataset( + dataset_path=run_cfg["dataset_path"], + split=split, + sampling_rate=run_cfg["sampling_rate"], + label_subset=run_cfg["label_subset"], + **run_cfg["data_kwargs"], + ) + + +def load_embeds( + embeds_dir: str, + split: str, + device: str, + run_cfg: dict, + chunk_rows: int = 500_000, +) -> torch.Tensor: + arr = np.load(os.path.join(embeds_dir, f"{split}_embeds.npy"), mmap_mode="r") + assert arr.shape[1] == run_cfg["n_cols_cached"], ( + f"{split} embeds have {arr.shape[1]} prototype columns, the embeds dir " + f"config implies {run_cfg['n_cols_cached']}" + ) + cols = run_cfg["keep_cols"] + shape = arr.shape if cols is None else (arr.shape[0], len(cols)) + print(f"loading {split} embeds {arr.shape} as {shape} to {device}") + out = torch.empty(shape, dtype=torch.float32, device=device) + for i in range(0, arr.shape[0], chunk_rows): + # subset on the host so dropped-branch columns never reach the GPU + chunk = np.array(arr[i : i + chunk_rows]) + if cols is not None: + chunk = chunk[:, cols] + out[i : i + chunk_rows].copy_(torch.from_numpy(chunk)) + return out + + +@torch.no_grad() +def compute_logits(W: torch.Tensor, X: torch.Tensor, chunk_rows: int = 65536): + return torch.cat( + [X[i : i + chunk_rows] @ W.T for i in range(0, len(X), chunk_rows)] + ) + + +@torch.no_grad() +def sparsity_stats(W: torch.Tensor, off_class_mask: torch.Tensor) -> dict: + abs_W = W.abs() + off_l1 = (abs_W * off_class_mask).sum() + off_abs = abs_W[off_class_mask.bool()] + return { + "off_class_l1": off_l1.item(), + "on_class_l1": (abs_W * (1 - off_class_mask)).sum().item(), + # share of total weight mass on off-class connections (lower = more interpretable) + "off_class_frac": (off_l1 / abs_W.sum()).item(), + # Adam + L1 subgradient does not produce exact zeros, so report near-zero + "off_class_near_zero_frac": (off_abs < 1e-3).float().mean().item(), + } + + +def train_head( + *, # enforce kwargs + params: dict, + X_train: torch.Tensor, + Y_train: torch.Tensor, + X_val: torch.Tensor, + Y_val: torch.Tensor, + assign: torch.Tensor, + pos_weight: torch.Tensor, + max_epochs: int, + early_stopping_patience: int, + lr_patience: int, + seed: int, + trial: optuna.Trial | None = None, +) -> tuple[torch.Tensor, dict]: + device = X_train.device + n_train, n_labels = Y_train.shape + off_class_mask = 1.0 - assign + lam_l1, batch_size = params["lam_l1"], params["batch_size"] + # exponent interpolates between unweighted (0) and full prevalence inversion + # (1, i.e. BaseTSDataset.get_label_weights as used by the Lightning stages) + pos_weight = pos_weight ** params["pos_weight_exp"] + + torch.manual_seed(seed) + gen = torch.Generator(device=device).manual_seed(seed) + # same masked init as ProtoECGNetFusion.cls + W = (assign - 0.5 * off_class_mask).clone().requires_grad_(True) # (L, R) + opt = torch.optim.AdamW([W], lr=params["lr"], weight_decay=params["weight_decay"]) + sched = torch.optim.lr_scheduler.ReduceLROnPlateau( + opt, mode="min", factor=0.1, patience=lr_patience + ) + + best_ap, best_epoch, best_W, best_metrics = -1.0, -1, None, {} + epoch = -1 + for epoch in trange(max_epochs, desc="Epoch"): + perm = torch.randperm(n_train, device=device, generator=gen) + train_loss_sum = torch.zeros((), device=device) + for start in trange(0, n_train, batch_size, desc="Step"): + idx = perm[start : start + batch_size] + logits = X_train[idx] @ W.T # (B, L) + cls_loss = F.binary_cross_entropy_with_logits( + logits, Y_train[idx], pos_weight=pos_weight, reduction="mean" + ) + penalty = (W.abs() * off_class_mask).sum() + loss = cls_loss + lam_l1 * penalty + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + train_loss_sum += loss.detach() * len(idx) + + with torch.no_grad(): + val_logits = compute_logits(W, X_val) + val_cls_loss = F.binary_cross_entropy_with_logits( + val_logits, Y_val, pos_weight=pos_weight, reduction="mean" + ).item() + val_penalty = (lam_l1 * (W.abs() * off_class_mask).sum()).item() + val_probs = val_logits.sigmoid() + val_ap = multilabel_average_precision( + val_probs, Y_val.int(), num_labels=n_labels, average="macro" + ).item() + val_auroc = multilabel_auroc( + val_probs, Y_val.int(), num_labels=n_labels, average="macro" + ).item() + val_loss = val_cls_loss + val_penalty + sched.step(val_loss) + print( + f" epoch {epoch:3d} | train_loss {(train_loss_sum / n_train).item():.4f} " + f"| val_loss {val_loss:.4f} (cls {val_cls_loss:.4f}) " + f"| val_ap {val_ap:.4f} | val_auroc {val_auroc:.4f} " + f"| lr {opt.param_groups[0]['lr']:.2e}", + flush=True, + ) + + if val_ap > best_ap: + best_ap, best_epoch = val_ap, epoch + best_W = W.detach().clone() + best_metrics = { + "val_ap": val_ap, + "val_auroc": val_auroc, + "val_classification_loss": val_cls_loss, + "val_loss": val_loss, + "best_epoch": epoch, + } + elif epoch - best_epoch >= early_stopping_patience: + break + + if trial is not None: + trial.report(val_ap, epoch) + if trial.should_prune(): + raise optuna.TrialPruned() + + assert best_W is not None + best_metrics["n_epochs"] = epoch + 1 + best_metrics.update(sparsity_stats(best_W, off_class_mask)) + return best_W, best_metrics + + +def tune(args, run_cfg: dict, study: optuna.Study, trials_dir: str): + device = args.device + assign = run_cfg["assign"].to(device) + + train_ds, val_ds = load_dataset(run_cfg, "train"), load_dataset(run_cfg, "val") + Y_train, Y_val = train_ds.labels, val_ds.labels + assert Y_train is not None and Y_val is not None + pos_weight = train_ds.get_label_weights().float() + + X_train = load_embeds(args.embeds_dir, "train", device, run_cfg) + X_val = load_embeds(args.embeds_dir, "val", device, run_cfg) + for X, Y in [(X_train, Y_train), (X_val, Y_val)]: + assert X.shape[0] == Y.shape[0], f"{X.shape} vs {Y.shape}" + assert X.shape[1] == assign.shape[1], f"{X.shape} vs {assign.shape}" + Y_train, Y_val = Y_train.float().to(device), Y_val.float().to(device) + pos_weight = pos_weight.to(device) + + def objective(trial: optuna.Trial) -> float: + params = { + "lam_l1": trial.suggest_float("lam_l1", 1e-8, 1e-2, log=True), + "lr": trial.suggest_float("lr", 1e-4, 3e-2, log=True), + "weight_decay": trial.suggest_float("weight_decay", 1e-6, 1e-1, log=True), + "batch_size": trial.suggest_categorical("batch_size", [512, 2048, 8192]), + "pos_weight_exp": trial.suggest_float("pos_weight_exp", 0.0, 1.0), + } + print(f"trial {trial.number}: {params}", flush=True) + W, metrics = train_head( + params=params, + X_train=X_train, + Y_train=Y_train, + X_val=X_val, + Y_val=Y_val, + assign=assign, + pos_weight=pos_weight, + max_epochs=args.max_epochs, + early_stopping_patience=args.early_stopping_patience, + lr_patience=args.lr_patience, + seed=args.seed + trial.number, + trial=trial, + ) + for k, v in metrics.items(): + trial.set_user_attr(k, v) + torch.save( + {"cls.weight": W.cpu()}, os.path.join(trials_dir, f"{trial.number}.pt") + ) + print(f"trial {trial.number} done: {metrics}", flush=True) + return metrics["val_ap"] + + if len(study.trials) == 0: + study.enqueue_trial(BASELINE_PARAMS) + n_finished = len( + [t for t in study.trials if t.state.is_finished()] + ) # resume toward n_trials total + study.optimize(objective, n_trials=max(args.n_trials - n_finished, 0)) + + +def export(args, run_cfg: dict, study: optuna.Study, trials_dir: str): + trials = study.trials_dataframe( + attrs=("number", "state", "value", "params", "user_attrs") + ) + trials = trials.rename( + columns=lambda c: c.removeprefix("params_").removeprefix("user_attrs_") + ) + trials = trials.sort_values("value", ascending=False) + trials.to_csv(os.path.join(args.output_dir, "trials.csv"), index=False) + + done = trials[trials["state"] == "COMPLETE"] + if len(done) == 0: + raise ValueError("No completed trials to select from") + best_ap = done["value"].max() + candidates = done[done["value"] >= best_ap - args.ap_tolerance] + selected = candidates.sort_values("off_class_frac").iloc[0] + summary_cols = [ + "number", "value", "val_auroc", "lam_l1", "lr", "weight_decay", + "batch_size", "pos_weight_exp", "off_class_frac", + "off_class_near_zero_frac", "best_epoch", + ] # fmt: skip + print(done[summary_cols].head(20).to_string(index=False)) + print( + f"selected trial {int(selected['number'])} (ap tolerance {args.ap_tolerance})" + ) + + W = torch.load(os.path.join(trials_dir, f"{int(selected['number'])}.pt"))[ + "cls.weight" + ] + W = W.to(args.device) + torch.save({"cls.weight": W.cpu()}, os.path.join(args.output_dir, "cls.pt")) + for split in ["val", "test"]: + X = load_embeds(args.embeds_dir, split, args.device, run_cfg) + assert X.shape[1] == W.shape[1], f"{X.shape} vs {W.shape}" + # column order is label_names, same layout as PredictionWriter probs + probs = compute_logits(W, X).sigmoid().cpu().numpy() + np.save(os.path.join(args.output_dir, f"{split}_probs.npy"), probs) + del X + + # the eval script reads the retained labels from here via --label-subset-config + labels_cfg = {"data": {"init_args": {"label_subset": run_cfg["label_names"]}}} + OmegaConf.save( + OmegaConf.create(labels_cfg), os.path.join(args.output_dir, "labels.yaml") + ) + + selection = { + "trial": int(selected["number"]), + "ap_tolerance": args.ap_tolerance, + "best_val_ap": float(best_ap), + "drop_branches": sorted(args.drop_branches), + "label_names": run_cfg["label_names"], + **{ + k: (v.item() if hasattr(v, "item") else v) + for k, v in selected.items() + if k not in {"number", "state"} + }, + } + with open(os.path.join(args.output_dir, "selected.json"), "w") as f: + json.dump(selection, f, indent=2, default=str) + print(f"exported selected head and probs to {args.output_dir}") + + +def tag_dropped_branches(study: optuna.Study, drop_branches: list[str], where: str): + """ + heads with different branches dropped have different shapes, so trials from one + ablation can neither be resumed nor selected against another - record the + ablation on the study and refuse to reuse a study that used a different one + """ + dropped = sorted(drop_branches) + # studies predating --drop-branches have trials but no attr, i.e. dropped nothing + prev = study.user_attrs.get("drop_branches", [] if len(study.trials) > 0 else None) + if prev is not None and sorted(prev) != dropped: + raise ValueError( + f"study in {where} was run with drop_branches={sorted(prev)}, " + f"not {dropped}; use a fresh --output-dir" + ) + study.set_user_attr("drop_branches", dropped) + + +def main(): + if os.environ.get("HIGH_MEMORY", None) is not None: + # this stage only needs the cached similarities, so materializing the + # waveform matrix would cost ~500 GB of RAM for nothing + raise ValueError("Do not set HIGH_MEMORY for this script") + args = parse_args() + os.makedirs(args.output_dir, exist_ok=True) + trials_dir = os.path.join(args.output_dir, "trials") + os.makedirs(trials_dir, exist_ok=True) + + run_cfg = load_run_config(args.embeds_dir, args.drop_branches) + study = optuna.create_study( + study_name="fusion-classifier", + storage=f"sqlite:///{os.path.join(args.output_dir, 'study.db')}", + load_if_exists=True, + direction="maximize", + sampler=optuna.samplers.TPESampler(seed=args.seed), + pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=5), + ) + tag_dropped_branches(study, args.drop_branches, args.output_dir) + if not args.export_only: + tune(args, run_cfg, study, trials_dir) + torch.cuda.empty_cache() # release train embeds before export + export(args, run_cfg, study, trials_dir) + + +if __name__ == "__main__": + main() diff --git a/large-user-study/scripts/queue-protoecgnet-heedb-overread.sh b/large-user-study/scripts/queue-protoecgnet-heedb-overread.sh new file mode 100755 index 0000000..fbbf156 --- /dev/null +++ b/large-user-study/scripts/queue-protoecgnet-heedb-overread.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +set -e + +# cache first so the arms don't race to build the same filtered caches +CACHE_JOB=$(sbatch --parsable run-cache-heedb-overread.sh) +echo "cache: $CACHE_JOB" + +ARM_JOBS=() +for arm in 1D-global 2D-global 2D-partial; do + job=$(sbatch --parsable --dependency=afterok:$CACHE_JOB run-protoecgnet-heedb-$arm.sh) + echo "$arm: $job" + ARM_JOBS+=($job) +done + +deps=$(IFS=:; echo "${ARM_JOBS[*]}") +FUSION_JOB=$(sbatch --parsable --dependency=afterok:$deps run-protoecgnet-heedb-fusion.sh) +echo "fusion: $FUSION_JOB" diff --git a/large-user-study/scripts/run-cache-heedb-overread.sh b/large-user-study/scripts/run-cache-heedb-overread.sh new file mode 100755 index 0000000..2828d55 --- /dev/null +++ b/large-user-study/scripts/run-cache-heedb-overread.sh @@ -0,0 +1,20 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=450gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protoecgnet-heedb-%j.out + +set -e + +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL/large-user-study + +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +python _cache_heedb_overread.py diff --git a/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh new file mode 100755 index 0000000..653077b --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh @@ -0,0 +1,74 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=450gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protoecgnet-heedb-%j.out + +set -e + +DATASET_PATH=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protoecgnet-heedb-overread +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL/large-user-study + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# experiment parameters +EXP_NAME="1D-global" + +# train via label supervised prototype learning +python -m protossl.protoecgnet_trainer \ + --pipeline-stage learn-prototypes-supervised \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +# project in the training dataset +python -m protossl.protoecgnet_trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python -m protossl.protoecgnet_trainer \ + --pipeline-stage train-classifier \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ +--label-subset-config $REPO_ROOT/configs/$EXP_NAME.yaml \ +--probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/test_probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/test_probs.npy $RUN_DIR/$EXP_NAME/test_probs.npy diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh new file mode 100755 index 0000000..58211e2 --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh @@ -0,0 +1,74 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=450gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protoecgnet-heedb-%j.out + +set -e + +DATASET_PATH=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protoecgnet-heedb-overread +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL/large-user-study + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# experiment parameters +EXP_NAME="2D-global" + +# train via label supervised prototype learning +python -m protossl.protoecgnet_trainer \ + --pipeline-stage learn-prototypes-supervised \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +# project in the training dataset +python -m protossl.protoecgnet_trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python -m protossl.protoecgnet_trainer \ + --pipeline-stage train-classifier \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ +--label-subset-config $REPO_ROOT/configs/$EXP_NAME.yaml \ +--probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/test_probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/test_probs.npy $RUN_DIR/$EXP_NAME/test_probs.npy diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh new file mode 100755 index 0000000..9f1fba5 --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh @@ -0,0 +1,74 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=450gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protoecgnet-heedb-%j.out + +set -e + +DATASET_PATH=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protoecgnet-heedb-overread +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL/large-user-study + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# experiment parameters +EXP_NAME="2D-partial" + +# train via label supervised prototype learning +python -m protossl.protoecgnet_trainer \ + --pipeline-stage learn-prototypes-supervised \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +# project in the training dataset +python -m protossl.protoecgnet_trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python -m protossl.protoecgnet_trainer \ + --pipeline-stage train-classifier \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ +--label-subset-config $REPO_ROOT/configs/$EXP_NAME.yaml \ +--probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/test_probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/test_probs.npy $RUN_DIR/$EXP_NAME/test_probs.npy diff --git a/large-user-study/scripts/run-protoecgnet-heedb-fusion-embed.sh b/large-user-study/scripts/run-protoecgnet-heedb-fusion-embed.sh new file mode 100755 index 0000000..b59b78d --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-fusion-embed.sh @@ -0,0 +1,53 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=500gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protoecgnet-heedb-%j.out + +set -e + +DATASET_PATH=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protoecgnet-heedb +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL/large-user-study + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# experiment parameters +EXP_NAME="fusion-tuned" + +python -m protossl.protoecgnet_trainer \ + --pipeline-stage compute-fusion-embeddings \ + --config $REPO_ROOT/configs/fusion.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --model.branches+=protossl.models._protoecgnet.BranchCfg \ + --model.branches.name=1D-global \ + --model.branches.config=$REPO_ROOT/configs/1D-global.yaml \ + --model.branches.pretrained_weights=$RUN_DIR/1D-global-tuned/project-prototypes-supervised/latest/proj.ckpt \ + --model.branches+=protossl.models._protoecgnet.BranchCfg \ + --model.branches.name=2D-global \ + --model.branches.config=$REPO_ROOT/configs/2D-global.yaml \ + --model.branches.pretrained_weights=$RUN_DIR/2D-global-tuned/project-prototypes-supervised/latest/proj.ckpt \ + --model.branches+=protossl.models._protoecgnet.BranchCfg \ + --model.branches.name=2D-partial \ + --model.branches.config=$REPO_ROOT/configs/2D-partial.yaml \ + --model.branches.pretrained_weights=$RUN_DIR/2D-partial-tuned/project-prototypes-supervised/latest/proj.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 diff --git a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh new file mode 100755 index 0000000..8a190d3 --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh @@ -0,0 +1,61 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=450gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protoecgnet-heedb-%j.out + +set -e + +DATASET_PATH=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protoecgnet-heedb-overread +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL/large-user-study + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# experiment parameters +EXP_NAME="fusion" + +python -m protossl.protoecgnet_trainer \ + --pipeline-stage train-fusion-classifier \ + --config $REPO_ROOT/configs/$EXP_NAME.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ + --model.branches+=protossl.models._protoecgnet.BranchCfg \ + --model.branches.name=1D-global \ + --model.branches.config=$REPO_ROOT/configs/1D-global.yaml \ + --model.branches.pretrained_weights=$RUN_DIR/1D-global/project-prototypes-supervised/latest/proj.ckpt \ + --model.branches+=protossl.models._protoecgnet.BranchCfg \ + --model.branches.name=2D-global \ + --model.branches.config=$REPO_ROOT/configs/2D-global.yaml \ + --model.branches.pretrained_weights=$RUN_DIR/2D-global/project-prototypes-supervised/latest/proj.ckpt \ + --model.branches+=protossl.models._protoecgnet.BranchCfg \ + --model.branches.name=2D-partial \ + --model.branches.config=$REPO_ROOT/configs/2D-partial.yaml \ + --model.branches.pretrained_weights=$RUN_DIR/2D-partial/project-prototypes-supervised/latest/proj.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ +--probs-npy $RUN_DIR/$EXP_NAME/train-fusion-classifier/latest/test_probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +cp $RUN_DIR/$EXP_NAME/train-fusion-classifier/latest/test_probs.npy $RUN_DIR/$EXP_NAME/test_probs.npy diff --git a/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh b/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh new file mode 100644 index 0000000..0aa061e --- /dev/null +++ b/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh @@ -0,0 +1,55 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=200gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protoecgnet-heedb-%j.out + +set -e + +DATASET_PATH=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protoecgnet-heedb +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL/large-user-study + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# experiment parameters +EXP_NAME="fusion-tuned" +# space separated branch names to ablate out of the head, empty keeps all of them. +# the cached embeddings always hold every branch, so this needs no re-embedding - +# but each setting gets its own output dir (and so its own optuna study) +DROP_BRANCHES="2D-global" +DROP_SUFFIX=${DROP_BRANCHES:+-drop-${DROP_BRANCHES// /-}} +EMBEDS_DIR=$RUN_DIR/$EXP_NAME/compute-fusion-embeddings/latest +OUTPUT_DIR=$RUN_DIR/$EXP_NAME/tune-fusion-classifier$DROP_SUFFIX + +# re-running resumes the optuna study (sqlite in $OUTPUT_DIR) up to --n-trials total +# use --export-only --ap-tolerance X to reselect without retraining +python _tune_fusion_classifier.py \ + --embeds-dir $EMBEDS_DIR \ + --output-dir $OUTPUT_DIR \ + ${DROP_BRANCHES:+--drop-branches $DROP_BRANCHES} \ + --n-trials 50 + +# labels.yaml is written by the tuning script: dropping a branch can drop labels, +# so the eval has to score the same label subset the head was tuned on +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--data-kwargs '{"heedb_split_type": "by-label"}' \ +--label-subset-config $OUTPUT_DIR/labels.yaml \ +--probs-npy $OUTPUT_DIR/test_probs.npy \ +--output-path $OUTPUT_DIR + +# suffixed so an ablation does not clobber the all-branches probs +cp $OUTPUT_DIR/test_probs.npy $RUN_DIR/$EXP_NAME/probs$DROP_SUFFIX.npy diff --git a/protossl/datasets/_base_dataset.py b/protossl/datasets/_base_dataset.py index c95f995..66ab984 100644 --- a/protossl/datasets/_base_dataset.py +++ b/protossl/datasets/_base_dataset.py @@ -19,6 +19,18 @@ shutil.rmtree(CACHE_DIR) +def get_cache_file( + *, # enforce kwargs + dataset_path: str, + split: SPLIT_T, + sampling_rate: int, +) -> str: + # use a short hash of the source dataest path to distinguish different sources (e.g. if using different dataset subsets) + identifier = f"{dataset_path.rstrip(os.sep)}_{split}_{sampling_rate}" + hashed = hashlib.md5(identifier.encode("utf-8")).hexdigest()[:8] + return os.path.join(CACHE_DIR, f"{hashed}.pt") + + def load_cached_data( *, # enforce kwargs load_transform_data_fn: Callable[[], torch.Tensor], @@ -28,11 +40,9 @@ def load_cached_data( ): # loading and processing can be time intensive, so cache transformed data if it doesn't already exist os.makedirs(CACHE_DIR, exist_ok=True) - - # use a short hash of the source dataest path to distinguish different sources (e.g. if using different dataset subsets) - identifier = f"{dataset_path.rstrip(os.sep)}_{split}_{sampling_rate}" - hashed = hashlib.md5(identifier.encode("utf-8")).hexdigest()[:8] - cache_file = os.path.join(CACHE_DIR, f"{hashed}.pt") + cache_file = get_cache_file( + dataset_path=dataset_path, split=split, sampling_rate=sampling_rate + ) print("=================load_cached_data==================") print(f"Dataset parameters: ({dataset_path}, {split}, {sampling_rate}Hz)") diff --git a/protossl/datasets/_heedb_dataset.py b/protossl/datasets/_heedb_dataset.py index 10f4927..68c7e62 100644 --- a/protossl/datasets/_heedb_dataset.py +++ b/protossl/datasets/_heedb_dataset.py @@ -1,5 +1,7 @@ +import hashlib import os from pathlib import Path +from typing import Literal import numpy as np import pandas as pd @@ -8,6 +10,7 @@ from tqdm import tqdm from ..defines import ( + CACHE_DIR, HEEDB_EUH_CLIPPED_MEANS, HEEDB_EUH_CLIPPED_STDS, HEEDB_EUH_LOWERS, @@ -21,9 +24,25 @@ SPLIT_T, STANDARD_LEAD_ORDER, ) -from ._base_dataset import BaseTSDataset, load_cached_data, validate_label_subset +from ._base_dataset import ( + BaseTSDataset, + get_cache_file, + load_cached_data, + validate_label_subset, +) from .streaming_loaders import StreamingECGWaveforms +# fmt: off +LABEL_SRC_MAPPING = { + "original_muse": ("diagnoses_acquisition.csv", "codes_software"), + "original_physician": ("diagnoses_acquisition.csv", "codes_physician"), + "new_muse": ("diagnoses_v24.csv", "codes"), +} +# fmt: on +LABEL_SRC_T = Literal["original_muse", "original_physician", "new_muse"] +HEEDB_SPLIT_T = Literal["by-year", "by-label"] +BY_LABEL_SPLIT_HASH = "373878c0ce8ff856a22b738dbbac4830476568b8c482b5d597ab9624ae3bacaf" # pragma: allowlist secret + heedb_lead_order = [l.lower() for l in HEEDB_LEAD_ORDER] standard_lead_order = [l.lower() for l in STANDARD_LEAD_ORDER] assert all([c == s for c, s in zip(heedb_lead_order, standard_lead_order)]) @@ -34,8 +53,8 @@ # global cache variables FULL_META = None -MGB_FNAME_TO_CODE = None -EUH_FNAME_TO_CODE = None +MGB_FNAME_TO_CODE: dict[str, dict[str, str]] = dict() +EUH_FNAME_TO_CODE: dict[str, dict[str, str]] = dict() class HeedbECGDataset(BaseTSDataset): @@ -46,13 +65,32 @@ def __init__( split: SPLIT_T, sampling_rate: int, label_subset: list[str] | None = None, + label_src: LABEL_SRC_T = "original_physician", + heedb_split_type: HEEDB_SPLIT_T = "by-year", + drop_blank_overread: bool = False, ): - df = get_heedb_metadata(dataset_path) + full_df = get_heedb_metadata(dataset_path, heedb_split_type=heedb_split_type) - df = df[df["split"] == split].reset_index(drop=True) + df = full_df[full_df["split"] == split].reset_index(drop=True) + has_overread = None + if drop_blank_overread: + # HEEDB v5 likely has some errors with missing physician overreads, + # so a blank physician overread can't be told apart from a physician + # deleting every statement. Thus, treat these rows as unlabeled + # rather than all-negative. filtered after split assignment so splits + # are unchanged and ecg_id (sample_id) stays stable + if label_src != "original_physician": + raise ValueError( + "drop_blank_overread only applies to label_src='original_physician', " + f"got: {label_src}" + ) + has_overread = get_heedb_overread_mask(dataset_path, df) + df = df[has_overread].reset_index(drop=True) self.source_ids = torch.as_tensor(df["patient_id"].to_numpy()) self.sample_ids = torch.as_tensor(df["ecg_id"].to_numpy()) - self.labels = torch.as_tensor(get_heedb_labels(dataset_path, df, label_subset)) + self.labels = torch.as_tensor( + get_heedb_labels(dataset_path, df, label_subset, label_src) + ) self._df = df streaming_ecgs = StreamingECGWaveforms( @@ -89,7 +127,35 @@ def __init__( print("WARNING:") print("===================================================") + base_split = ( + split + if heedb_split_type == "by-year" + else f"{split}-{heedb_split_type}" + ) + unfiltered_cache = get_cache_file( + dataset_path=dataset_path, + split=base_split, # type: ignore + sampling_rate=sampling_rate, + ) + cache_split = ( + base_split + if not drop_blank_overread + else f"{base_split}-drop-blank-overread" + ) + def load_transform_data_fn() -> torch.Tensor: + if drop_blank_overread and os.path.exists(unfiltered_cache): + # rows of the unfiltered cache are the unfiltered split df in order, so + # index it with the overread mask instead of re-streaming every file. + # mmap keeps peak memory at roughly the size of the filtered tensor + print(f"deriving filtered cache from: {unfiltered_cache}") + X = torch.load(unfiltered_cache, mmap=True) + assert has_overread is not None + assert X.shape[0] == len(has_overread), ( + f"unfiltered cache has {X.shape[0]} rows, " + f"expected {len(has_overread)}" + ) + return X[torch.from_numpy(has_overread)] # fmt: off print("WARNING:") print("WARNING: ABOUT TO LOAD ENTIRE HEEDB WAVEFORM MATRIX INTO MEMORY TO CACHE") @@ -111,7 +177,7 @@ def load_transform_data_fn() -> torch.Tensor: self.waveforms = load_cached_data( load_transform_data_fn=load_transform_data_fn, dataset_path=dataset_path, - split=split, + split=cache_split, # type: ignore sampling_rate=sampling_rate, ) @@ -120,89 +186,121 @@ def load_transform_data_fn() -> torch.Tensor: assert self.source_ids.shape[0] == self.labels.shape[0] -def get_heedb_metadata(heedb_path: str) -> pd.DataFrame: +def get_heedb_metadata( + heedb_path: str, + *, # enforce kwargs + heedb_split_type: HEEDB_SPLIT_T = "by-year", +) -> pd.DataFrame: _path = Path(heedb_path) global FULL_META if FULL_META is not None: return FULL_META.copy() print("================get_heedb_metadata=================") - print("reading HEEDB metadata") + print(f"using {heedb_split_type} splits") - # read harvard data - mgb = pd.read_csv( - _path / "I0001/metadata/metadata.csv", - usecols=["BDSPPatientID", "SexDSC", "AgeAtAcquisition", "FileName"], - ) - mgb["AgeAtAcquisition"] = mgb["AgeAtAcquisition"] / 365.2425 - mgb = mgb[(mgb["AgeAtAcquisition"] >= 18) & (mgb["SexDSC"].notna())] - mgb = mgb.rename( - columns={ - "BDSPPatientID": "patient_id", - "SexDSC": "sex", - "AgeAtAcquisition": "age", - "FileName": "fpath", - } - ) - print("read MGB data") + identifier = f"HEEDB_{heedb_path.rstrip(os.sep)}_{heedb_split_type}" + hashed = hashlib.md5(identifier.encode("utf-8")).hexdigest()[:8] + cache_file = os.path.join(CACHE_DIR, f"{hashed}.csv") + if os.path.exists(cache_file): + print(f"reading HEEDB metadata from on-disk cache: {cache_file}") + df = pd.read_csv(cache_file) + else: + print(f"reading HEEDB metadata from source: {heedb_path}") - # read emory data, slight differences - emory = pd.read_csv( - _path / "I0006/metadata/metadata.csv", - usecols=["BDSPPatientID", "Sex", "AgeAtAcquisition", "FileName"], - ) - emory["AgeAtAcquisition"] = emory["AgeAtAcquisition"] / 365.2425 - emory = emory[ - (emory["AgeAtAcquisition"] >= 18) - & (emory["Sex"].notna()) - & (emory["BDSPPatientID"].notna()) - ] - emory = emory.rename( - columns={ - "BDSPPatientID": "patient_id", - "Sex": "sex", - "AgeAtAcquisition": "age", - "FileName": "fpath", - } - ) - assert (emory["patient_id"].astype(int) == emory["patient_id"]).all() - emory["patient_id"] = emory["patient_id"].astype(int) - - # bad files on emory side - emory_exclude = {"WFDB/2013/MUSE_20200225_081000_06000"} - emory = emory[~emory["fpath"].isin(emory_exclude)] - - print("read Emory data") - - # join together - mgb["source"] = "mgb" - emory["source"] = "emory" - assert len(set(mgb["patient_id"]) & set(emory["patient_id"])) == 0 - df = pd.concat([mgb, emory], ignore_index=True) # MGB, then EUH - df.index.name = "ecg_id" - df["year"] = df["fpath"].str[1:].str.split("/").str[1].astype(int) - df = df.reset_index()[ - ["ecg_id", "patient_id", "age", "sex", "year", "source", "fpath"] - ] - - # emory data ends in 2018 so val/test are all MGB data - df["split"] = "no-split" - df.loc[~df["year"].isin([2021, 2022]), "split"] = "train" - df.loc[df["year"] == 2021, "split"] = "val" - df.loc[df["year"] == 2022, "split"] = "test" - - full_paths = [] - for f, src in zip(df["fpath"], df["source"]): - if src == "mgb": - # harvard paths start with "/S...", the level just under WFDB - p = _path / "I0001/WFDB" / f[1:] - elif src == "emory": - # emory paths start with "WFDB/..." - p = _path / "I0006" / f + # read harvard data + mgb = pd.read_csv( + _path / "I0001/metadata/metadata.csv", + usecols=["BDSPPatientID", "SexDSC", "AgeAtAcquisition", "FileName"], + ) + mgb["AgeAtAcquisition"] = mgb["AgeAtAcquisition"] / 365.2425 + mgb = mgb[(mgb["AgeAtAcquisition"] >= 18) & (mgb["SexDSC"].notna())] + mgb = mgb.rename( + columns={ + "BDSPPatientID": "patient_id", + "SexDSC": "sex", + "AgeAtAcquisition": "age", + "FileName": "fpath", + } + ) + print("read MGB data") + + # read emory data, slight differences + emory = pd.read_csv( + _path / "I0006/metadata/metadata.csv", + usecols=["BDSPPatientID", "Sex", "AgeAtAcquisition", "FileName"], + ) + emory["AgeAtAcquisition"] = emory["AgeAtAcquisition"] / 365.2425 + emory = emory[ + (emory["AgeAtAcquisition"] >= 18) + & (emory["Sex"].notna()) + & (emory["BDSPPatientID"].notna()) + ] + emory = emory.rename( + columns={ + "BDSPPatientID": "patient_id", + "Sex": "sex", + "AgeAtAcquisition": "age", + "FileName": "fpath", + } + ) + assert (emory["patient_id"].astype(int) == emory["patient_id"]).all() + emory["patient_id"] = emory["patient_id"].astype(int) + + # bad files on emory side + emory_exclude = {"WFDB/2013/MUSE_20200225_081000_06000"} + emory = emory[~emory["fpath"].isin(emory_exclude)] + print("read Emory data") + + # join together + mgb["source"] = "mgb" + emory["source"] = "emory" + assert len(set(mgb["patient_id"]) & set(emory["patient_id"])) == 0 + df = pd.concat([mgb, emory], ignore_index=True) # MGB, then EUH + df.index.name = "ecg_id" + df["year"] = df["fpath"].str[1:].str.split("/").str[1].astype(int) + df = df.reset_index()[ + ["ecg_id", "patient_id", "age", "sex", "year", "source", "fpath"] + ] + + df["split"] = "train" + if heedb_split_type == "by-year": + # emory data ends in 2018 so val/test are all MGB data + df.loc[df["year"] == 2021, "split"] = "val" + df.loc[df["year"] == 2022, "split"] = "test" + elif heedb_split_type == "by-label": + # law of large numbers, we just generate random splits and it's close enough + # to splits which preserve independent label prevalence (need to check about cooccurrence) + n = int(len(df) * 0.05) + rng = np.random.default_rng(42) + val_test_idxs = rng.choice(len(df), size=n * 2, replace=False) + digest = hashlib.sha256(val_test_idxs.tobytes()).hexdigest() + if digest != BY_LABEL_SPLIT_HASH: + raise ValueError( + "Randomly generated split for preserving label prevalence has changed!" + ) + + df.loc[val_test_idxs[:n], "split"] = "val" + df.loc[val_test_idxs[n:], "split"] = "test" else: - raise ValueError(f"Unknown path structure for institution: {src}") - full_paths.append(p) - df["full_path"] = full_paths + raise ValueError( + f"Unknown how to create splits for HEEDB split type: {heedb_split_type}" + ) + + full_paths = [] + for f, src in zip(df["fpath"], df["source"]): + if src == "mgb": + # harvard paths start with "/S...", the level just under WFDB + p = _path / "I0001/WFDB" / f[1:] + elif src == "emory": + # emory paths start with "WFDB/..." + p = _path / "I0006" / f + else: + raise ValueError(f"Unknown path structure for institution: {src}") + full_paths.append(p) + df["full_path"] = full_paths + df.to_csv(cache_file, index=False) + print(f"saved HEEDB metadata to on-disk cache: {cache_file}") FULL_META = df.copy() print("===================================================") @@ -213,75 +311,185 @@ def get_heedb_labels( heedb_path: str, meta: pd.DataFrame, label_subset: list[str] | None = None, + label_src: LABEL_SRC_T = "original_physician", ) -> np.ndarray: print("=================make_heedb_labels=================") targets = HEEDB_TARGETS if label_subset is not None: validate_label_subset(label_subset, list(HEEDB_TARGETS)) targets = {label: HEEDB_TARGETS[label] for label in label_subset} - global MGB_FNAME_TO_CODE - global EUH_FNAME_TO_CODE - - def make_fname_to_code(institution) -> dict[str, str]: - if institution == "mgb": - subdir = "I0001" - elif institution == "emory": - subdir = "I0006" - else: - raise ValueError(f"Unknown subdir for institution: {institution}") - df = pd.read_csv( - os.path.join(heedb_path, subdir, "12SL_diagnoses/diagnoses_acquisition.csv") - ) - return { - fname: code_str - for fname, code_str in zip( - tqdm( - df["FileName"], desc=f"Creating {institution} fname to code mapping" - ), - df["codes_physician"], + + meta_hash = hash_path_list(meta["fpath"]) + targets_hash = hashlib.md5("\0".join(targets).encode("utf-8")).hexdigest()[:8] + identifier = f"HEEDB_labels_{heedb_path.rstrip(os.sep)}_{meta_hash}_{label_src}_{targets_hash}" + hashed = hashlib.md5(identifier.encode("utf-8")).hexdigest()[:8] + cache_file = os.path.join(CACHE_DIR, f"{hashed}.npy") + + if os.path.exists(cache_file): + print(f"reading HEEDB labels from on-disk cache: {cache_file}") + loaded = np.load(cache_file, allow_pickle=True).item() + data = loaded["data"] + n_ecgs, n_annots = loaded["n_ecgs"], loaded["n_annots"] + n_matched = loaded["n_matched"] + print("replaying below stats from cache:") + else: + print(f"reading HEEDB labels from source: {heedb_path}") + load_fname_to_code(heedb_path, label_src) # caches in global vars + code_to_label = {c: k for k, cs in targets.items() for c in cs} + label_to_idx = {k: i for i, k in enumerate(targets)} + + data = np.zeros((len(meta), len(targets)), dtype=np.long) + count = 0 + for meta_idx, (fname, institution) in enumerate( + zip( + tqdm(meta["fpath"], desc="Converting code string to labels"), + meta["source"], ) - } - - if MGB_FNAME_TO_CODE is None: - MGB_FNAME_TO_CODE = make_fname_to_code("mgb") - if EUH_FNAME_TO_CODE is None: - EUH_FNAME_TO_CODE = make_fname_to_code("emory") - code_to_label = {c: k for k, cs in targets.items() for c in cs} - label_to_idx = {k: i for i, k in enumerate(targets)} - - data = np.zeros((len(meta), len(targets)), dtype=np.long) - count = 0 - for meta_idx, (fname, institution) in enumerate( - zip( - tqdm(meta["fpath"], desc="Converting code string to labels"), meta["source"] - ) - ): - if institution == "mgb": - codes = MGB_FNAME_TO_CODE.get(fname, "MISSING") - elif institution == "emory": - codes = EUH_FNAME_TO_CODE.get(fname, "MISSING") - else: - raise ValueError(f"Unknown institution: {institution}") - if codes == "MISSING": - # separate missing vs empty (below) - continue - count += 1 - if isinstance(codes, float) and np.isnan(codes): - continue - assert isinstance(codes, str) - for code in codes.split(","): - # convert 12SL code to label - label = code_to_label.get(int(code), -1) - if label == -1: + ): + if institution == "mgb": + codes = MGB_FNAME_TO_CODE[label_src].get(fname, "MISSING") + elif institution == "emory": + codes = EUH_FNAME_TO_CODE[label_src].get(fname, "MISSING") + else: + raise ValueError(f"Unknown institution: {institution}") + if codes == "MISSING": + # separate missing vs empty (below) continue - # convert label to column idx - label_idx = label_to_idx.get(label, -1) # type: ignore - if label_idx == -1: + count += 1 + if isinstance(codes, float) and np.isnan(codes): continue - data[meta_idx, label_idx] = 1 + assert isinstance(codes, str) + for code in codes.split(","): + # convert 12SL code to label + label = code_to_label.get(int(code), -1) + if label == -1: + continue + # convert label to column idx + label_idx = label_to_idx.get(label, -1) # type: ignore + if label_idx == -1: + continue + data[meta_idx, label_idx] = 1 + + n_ecgs, n_matched = len(meta), count + n_annots = len(MGB_FNAME_TO_CODE[label_src]) + len(EUH_FNAME_TO_CODE[label_src]) + np.save( + cache_file, + { # type: ignore + "data": data, + "n_ecgs": n_ecgs, + "n_annots": n_annots, + "n_matched": n_matched, + }, + ) + print(f"saved HEEDB labels to on-disk cache: {cache_file}") print( - f"Of {len(meta)} ECGs and {len(MGB_FNAME_TO_CODE) + len(EUH_FNAME_TO_CODE)} annotations, " - f"{count} matched ({len(meta)-count} ECG annotations were missing and filled with 0s)" + f"Of {n_ecgs} ECGs and {n_annots} annotations, {n_matched} matched " + f"({n_ecgs-n_matched} ECG annotations were missing and filled with 0s)" ) print("===================================================") return data + + +def make_fname_to_code( + heedb_path: str, + institution: str, + label_csv: str, + label_col: str, +) -> dict[str, str]: + if institution == "mgb": + subdir = "I0001" + elif institution == "emory": + subdir = "I0006" + else: + raise ValueError(f"Unknown subdir for institution: {institution}") + df = pd.read_csv(os.path.join(heedb_path, subdir, "12SL_diagnoses", label_csv)) + return { + fname: code_str + for fname, code_str in zip( + tqdm( + # v24 file names are prefixed with '.' and suffixed with '.hea\n' + df["FileName"].str.strip(".").str.strip(".hea\n"), + desc=f"Creating {institution} fname to code mapping", + ), + df[label_col], + ) + } + + +def load_fname_to_code(heedb_path: str, label_src: LABEL_SRC_T): + # populates the global per-institution caches for this label source + _label_csv, _label_col = LABEL_SRC_MAPPING[label_src] + if label_src not in MGB_FNAME_TO_CODE: + MGB_FNAME_TO_CODE[label_src] = make_fname_to_code( + heedb_path, "mgb", _label_csv, _label_col + ) + if label_src not in EUH_FNAME_TO_CODE: + EUH_FNAME_TO_CODE[label_src] = make_fname_to_code( + heedb_path, "emory", _label_csv, _label_col + ) + + +def get_heedb_overread_mask(heedb_path: str, meta: pd.DataFrame) -> np.ndarray: + """ + Boolean mask over `meta` rows, True where the ECG has a non-blank physician + overread (`codes_physician` in diagnoses_acquisition.csv). Rows whose file isn't + in the diagnoses table at all are also False, but are counted separately. + """ + print("==============get_heedb_overread_mask==============") + label_src: LABEL_SRC_T = "original_physician" + meta_hash = hash_path_list(meta["fpath"]) + identifier = ( + f"HEEDB_overread_mask_{heedb_path.rstrip(os.sep)}_{meta_hash}_{label_src}" + ) + hashed = hashlib.md5(identifier.encode("utf-8")).hexdigest()[:8] + cache_file = os.path.join(CACHE_DIR, f"{hashed}.npy") + + if os.path.exists(cache_file): + print(f"reading HEEDB overread mask from on-disk cache: {cache_file}") + loaded = np.load(cache_file, allow_pickle=True).item() + mask = loaded["mask"] + n_blank, n_missing = loaded["n_blank"], loaded["n_missing"] + print("replaying below stats from cache:") + else: + print(f"reading HEEDB physician overreads from source: {heedb_path}") + load_fname_to_code(heedb_path, label_src) + mask = np.zeros(len(meta), dtype=bool) + n_blank, n_missing = 0, 0 + for meta_idx, (fname, institution) in enumerate( + zip( + tqdm(meta["fpath"], desc="Checking physician overreads"), meta["source"] + ) + ): + if institution == "mgb": + codes = MGB_FNAME_TO_CODE[label_src].get(fname, "MISSING") + elif institution == "emory": + codes = EUH_FNAME_TO_CODE[label_src].get(fname, "MISSING") + else: + raise ValueError(f"Unknown institution: {institution}") + if codes == "MISSING": + n_missing += 1 + elif not isinstance(codes, str) or codes.strip() == "": + # pandas reads an empty field as NaN + n_blank += 1 + else: + mask[meta_idx] = True + os.makedirs(CACHE_DIR, exist_ok=True) + np.save( + cache_file, + {"mask": mask, "n_blank": n_blank, "n_missing": n_missing}, # type: ignore + ) + print(f"saved HEEDB overread mask to on-disk cache: {cache_file}") + print( + f"Of {len(meta)} ECGs, keeping {mask.sum()} with a physician overread, dropping " + f"{n_blank} with a blank overread and {n_missing} not found in the diagnoses table" + ) + print("===================================================") + return mask + + +def hash_path_list(paths: pd.Series, chunk=100_000) -> str: + m = hashlib.sha256() + for i in range(0, len(paths), chunk): + m.update("\0".join(paths.iloc[i : i + chunk]).encode()) + m.update(b"\0") + return m.hexdigest() diff --git a/protossl/datasets/streaming_loaders/_streaming_base.py b/protossl/datasets/streaming_loaders/_streaming_base.py index 4b3d792..ab64084 100644 --- a/protossl/datasets/streaming_loaders/_streaming_base.py +++ b/protossl/datasets/streaming_loaders/_streaming_base.py @@ -13,3 +13,6 @@ def __getitem__(self, i: int) -> torch.Tensor: @abstractmethod def shape(self) -> tuple[int, ...]: pass + + def __len__(self) -> int: + return self.shape[0] diff --git a/protossl/defines.py b/protossl/defines.py index 7a8a2a3..b2e7858 100644 --- a/protossl/defines.py +++ b/protossl/defines.py @@ -11,6 +11,9 @@ "project-prototypes-supervised", "compute-embeddings", "train-classifier", + "train-fusion-classifier", + "compute-fusion-embeddings", + "train-contraster", ] LABEL_T = Literal["binary-multilabel", "multiclass"] @@ -80,7 +83,7 @@ ] CACHE_DIR = os.environ.get( - "CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache/pass_pclr_cache") + "CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache/protossl_cache") ) # EchoNext per-lead stats not defined because data comes prenormalized @@ -523,6 +526,162 @@ "WITH 1ST DEGREE AV BLOCK": [101], "WITH SINUS ARRHYTHMIA": [251], } +if os.environ.get("USE_HEEDB_150", False): + print("====================defines.py=====================") + print("USING HEEDB 150 LABEL SET") + print("===================================================") + HEEDB_TARGETS = { + "ABNORMAL ECG": [1699], + "NORMAL SINUS RHYTHM": [22], + "NORMAL ECG": [1684], + "SINUS RHYTHM": [19], + "SINUS BRADYCARDIA": [21], + "ATRIAL FIBRILLATION": [161], + "SINUS TACHYCARDIA": [23], + "OTHERWISE NORMAL ECG": [1687], + "LEFT AXIS DEVIATION": [372], + "PREMATURE VENTRICULAR COMPLEXES": [231], + "BORDERLINE ECG": [1693], + "RIGHT BUNDLE BRANCH BLOCK": [440], + "SEPTAL INFARCT": [700], + "LEFT ATRIAL ENLARGEMENT": [360], + "NONSPECIFIC T WAVE ABNORMALITY": [1140], + "LOW VOLTAGE QRS": [410], + "PREMATURE ATRIAL COMPLEXES": [222], + "ANTERIOR INFARCT": [740], + "INCOMPLETE RIGHT BUNDLE BRANCH BLOCK": [445], + "PREMATURE SUPRAVENTRICULAR COMPLEXES": [221], + "LEFT BUNDLE BRANCH BLOCK": [460], + "NONSPECIFIC T WAVE ABNORMALITY NOW EVIDENT IN": [1215], + "NONSPECIFIC T WAVE ABNORMALITY NO LONGER EVIDENT IN": [1214], + "T WAVE INVERSION NOW EVIDENT IN": [1220], + "LATERAL INFARCT": [760], + "NONSPECIFIC ST ABNORMALITY": [900, 1021], + "LEFT VENTRICULAR HYPERTROPHY": [541], + "T WAVE INVERSION NO LONGER EVIDENT IN": [1224], + "WITH RAPID VENTRICULAR RESPONSE": [171], + "QT HAS SHORTENED": [1251], + "QT HAS LENGTHENED": [1250], + "FUSION COMPLEXES": [244], + "ATRIAL FLUTTER": [162, 273, 288], + "MARKED SINUS BRADYCARDIA": [24], + "WITH SINUS ARRHYTHMIA": [251], + "NONSPECIFIC ST AND T WAVE ABNORMALITY": [1141], + "LEFT ANTERIOR FASCICULAR BLOCK": [470], + "RIGHT AXIS DEVIATION": [383], + "ECTOPIC ATRIAL RHYTHM": [64], + "UNDETERMINED RHYTHM": [299], + "ANTEROSEPTAL INFARCT": [810], + "RIGHTWARD AXIS": [380], + "ST NOW DEPRESSED IN": [1107], + "WITH SHORT PR": [102], + "WITH MARKED SINUS ARRHYTHMIA": [252], + "ST NO LONGER DEPRESSED IN": [1104], + "INVERTED T WAVES HAVE REPLACED NONSPECIFIC T WAVE ABNORMALITY IN": [1219], + "NON-SPECIFIC CHANGE IN ST SEGMENT IN": [1116, 1117], + "NONSPECIFIC T WAVE ABNORMALITY HAS REPLACED INVERTED T WAVES IN": [1218], + "JUNCTIONAL RHYTHM": [267], + "ELECTRONIC ATRIAL PACEMAKER": [292], + "ABERRANT CONDUCTION": [246], + "ELECTRONIC VENTRICULAR PACEMAKER": [290], + "T WAVE INVERSION LESS EVIDENT IN": [1223], + "ANTEROLATERAL INFARCT": [820], + "WITH REPOLARIZATION ABNORMALITY": [544], + "RSR' OR QR PATTERN IN V1 SUGGESTS RIGHT VENTRICULAR CONDUCTION DELAY": [450], + "T WAVE INVERSION MORE EVIDENT IN": [1221], + "WIDE QRS RHYTHM": [237], + "WITH PREMATURE VENTRICULAR OR ABERRANTLY CONDUCTED COMPLEXES": [173, 181], + "RIGHT ATRIAL ENLARGEMENT": [350], + "INFERIOR INFARCT": [780], + "INCOMPLETE LEFT BUNDLE BRANCH BLOCK": [465], + "VOLTAGE CRITERIA FOR LEFT VENTRICULAR HYPERTROPHY": [540], + "OR DIGITALIS EFFECT": [1671], + "BIFASCICULAR BLOCK": [480], + "ST NO LONGER ELEVATED IN": [1123], + "WITH SLOW VENTRICULAR RESPONSE": [172], + "ST ELEVATION NOW PRESENT IN": [1122], + "PREMATURE ECTOPIC COMPLEXES": [241], + "LEFT POSTERIOR FASCICULAR BLOCK": [471], + "T WAVE AMPLITUDE HAS DECREASED IN": [1203], + "WITH A COMPETING JUNCTIONAL PACEMAKER": [174], + "RIGHT SUPERIOR AXIS DEVIATION": [384], + "BIATRIAL ENLARGEMENT": [369], + "VENTRICULAR-PACED RHYTHM": [296], + "ATRIAL-PACED RHYTHM": [295], + "T WAVE AMPLITUDE HAS INCREASED IN": [1201], + "WITH QRS WIDENING": [543], + "WITH 1ST DEGREE AV BLOCK": [101], + "PROLONGED QT": [1143], + "WITH PROLONGED AV CONDUCTION": [190], + "RIGHT VENTRICULAR HYPERTROPHY": [520], + "WITH QRS WIDENING AND REPOLARIZATION ABNORMALITY": [545], + "ATRIAL-SENSED VENTRICULAR-PACED RHYTHM": [297], + "AV SEQUENTIAL OR DUAL CHAMBER ELECTRONIC PACEMAKER": [293], + "PULMONARY DISEASE PATTERN": [411], + "ACUTE MI / STEMI": [821], + "INFERIOR-POSTERIOR INFARCT": [801], + "NONSPECIFIC INTRAVENTRICULAR CONDUCTION DELAY": [487], + "PREMATURE VENTRICULAR AND FUSION COMPLEXES": [232], + "IN A PATTERN OF BIGEMINY": [234], + "AV DUAL-PACED RHYTHM": [298], + "SUPRAVENTRICULAR TACHYCARDIA": [266, 271], + "VENTRICULAR-PACED COMPLEXES": [184], + "WIDE QRS TACHYCARDIA": [235], + "RSR' PATTERN IN V1": [451], + "ST LESS DEPRESSED IN": [1105], + "VENTRICULAR TACHYCARDIA": [248, 272, 277], + "EARLY REPOLARIZATION": [1000], + "ST MORE DEPRESSED IN": [1106], + "ANTEROLATERAL LEADS": [1456], + "ELECTRONIC DEMAND PACING": [294], + "RBBB AND LEFT ANTERIOR FASCICULAR BLOCK": [478], + "LATERAL INJURY PATTERN": [940], + "BIVENTRICULAR PACEMAKER DETECTED": [289], + "SUSPECT UNSPECIFIED PACEMAKER FAILURE": [1669], + "WOLFF-PARKINSON-WHITE": [304], + "WITH VENTRICULAR ESCAPE COMPLEXES": [243], + "INFERIOR INJURY PATTERN": [950], + "CONSIDER RIGHT VENTRICULAR INVOLVEMENT IN ACUTE INFERIOR INFARCT": [806], + "ST ELEVATION HAS REPLACED ST DEPRESSION IN": [1124], + "NONSPECIFIC INTRAVENTRICULAR BLOCK": [482], + "MASKED BY FASCICULAR BLOCK": [782], + "PEDIATRIC ECG ANALYSIS": [2], + "BLOCKED": [32], + "WITH UNDETERMINED RHYTHM IRREGULARITY": [175], + "LEFTWARD AXIS": [370], + "WITH 2ND DEGREE SA BLOCK MOBITZ I": [112], + "ACUTE": [1423], + "ABNORMAL LEFT AXIS DEVIATION": [371], + "WITH COMPLETE HEART BLOCK": [106], + "NO P-WAVES FOUND": [31], + "ST LESS ELEVATED IN": [1121], + "WITH RETROGRADE CONDUCTION": [245], + "ST MORE ELEVATED IN": [1120], + "JUNCTIONAL BRADYCARDIA": [34], + "WITH VARIABLE AV BLOCK": [107], + "ANTERIOR INJURY PATTERN": [930], + "WITH JUNCTIONAL ESCAPE COMPLEXES": [242], + "ACUTE MI": [829], + "ACUTE PERICARDITIS": [901], + "POSTERIOR INFARCT": [802], + "IDIOVENTRICULAR RHYTHM": [238], + "WITH 2ND DEGREE SA BLOCK MOBITZ II": [111], + "R IN AVL": [530], + "SINUS/ATRIAL CAPTURE": [247], + "AV DUAL-PACED COMPLEXES": [185], + "INFEROLATERAL INJURY PATTERN": [962], + "RBBB AND LEFT POSTERIOR FASCICULAR BLOCK": [479], + "ANTEROLATERAL INJURY PATTERN": [961], + "ATRIAL-PACED COMPLEXES": [183], + "WITH SINUS PAUSE": [113], + "BIVENTRICULAR HYPERTROPHY": [570], + "ABNORMAL RIGHT AXIS DEVIATION": [381], + "SUPRAVENTRICULAR COMPLEXES": [188], + "WITH 2ND DEGREE AV BLOCK MOBITZ I": [103], + "WITH 2:1 AV CONDUCTION": [141], + "WITH AV DISSOCIATION": [108], + "MULTIFOCAL ATRIAL TACHYCARDIA": [280], + } HEEDB_LEAD_ORDER = [ "I", "II", diff --git a/protossl/lightning_utils.py b/protossl/lightning_utils.py index f79f9d3..c4112b3 100644 --- a/protossl/lightning_utils.py +++ b/protossl/lightning_utils.py @@ -7,7 +7,7 @@ from lightning.pytorch.loggers import WandbLogger from lightning.pytorch.utilities import rank_zero_only -from wandb.util import generate_id +from wandb.sdk.lib.runid import generate_id from .defines import STAGE_T diff --git a/protossl/models/__init__.py b/protossl/models/__init__.py index 098c963..8e34c37 100644 --- a/protossl/models/__init__.py +++ b/protossl/models/__init__.py @@ -1,5 +1,6 @@ from ._base_classifier import BaseClassifier from ._blackbox_classifier import BlackboxClassifier +from ._blackbox_contraster import BlackboxContraster from ._prototype_assigner import PrototypeAssigner from ._prototype_classifier import PrototypeClassifier from ._prototype_contraster import PrototypeContraster diff --git a/protossl/models/_blackbox_contraster.py b/protossl/models/_blackbox_contraster.py new file mode 100644 index 0000000..357dee8 --- /dev/null +++ b/protossl/models/_blackbox_contraster.py @@ -0,0 +1,143 @@ +from typing import get_args + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ..defines import BACKBONE_T, CONTRASTIVE_T, CONV_T +from ._pretrained_utils import PretrainedMixin +from .encoders import ResNet1D, ResNet2D + + +class BlackboxContraster(PretrainedMixin, nn.Module): + def __init__( + self, + *, # enforce kwargs + contrastive_pair_mode: CONTRASTIVE_T, + backbone_type: BACKBONE_T, + conv_type: CONV_T, + proj_dim: int | None = None, + init_log_temp: float = 0.07, + learnable_temp: bool = True, + pretrained_weights: str | None = None, + input_channels: int = 12, + # specified via model_kawrgs in trainer + ): + super().__init__() + + if contrastive_pair_mode not in get_args(CONTRASTIVE_T): + raise ValueError(f"Unknown contrastive_pair_mode={contrastive_pair_mode}.") + if contrastive_pair_mode != "pclr": + raise NotImplementedError( + f"Non-PCLR contrastive_pair_mode ({contrastive_pair_mode}) not yet supported for BlackboxContraster" + ) + if not backbone_type.startswith("resnet"): + raise NotImplementedError( + f"Non-resnet (lower-cased) backbone_type ({backbone_type}) not yet supported for BlackboxContraster" + ) + + if conv_type == "1D": + backbone_cls = ResNet1D + elif conv_type == "2D": + backbone_cls = ResNet2D + elif conv_type == "PANNS": + raise NotImplemented( + f"PANNS conv_type not yet supported for BlackboxContraster" + ) + else: + raise ValueError(f"Unknown conv_type={conv_type}") + + self.encoder = backbone_cls( + backbone_type=backbone_type, + input_channels=input_channels, + ) + emb_dim = self.encoder.emb_dim + if proj_dim is None: + proj_dim = emb_dim // 2 + self.proj = nn.Sequential( + nn.Linear(emb_dim, emb_dim), + nn.ReLU(inplace=True), + nn.Linear(emb_dim, proj_dim), + ) + + self.log_temperature = nn.Parameter( + torch.ones([]) * np.log(1 / init_log_temp), + requires_grad=learnable_temp, + ) + self.contrastive_pair_mode = contrastive_pair_mode + if pretrained_weights is not None: + self.load_pretrained_weights(pretrained_weights) + + def _project_pair( + self, x1: torch.Tensor, x2: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + assert x1.shape == x2.shape + + x1 = self.encoder(x1) + x2 = self.encoder(x2) + + x1 = self.proj(x1) + x2 = self.proj(x2) + + return x1, x2 + + def forward( + self, + x1: torch.Tensor, + x2: torch.Tensor, + ) -> dict[str, torch.Tensor]: + z1, z2 = self._project_pair(x1, x2) + simclr_loss = self._simclr_loss(z1, z2) + + return {"SimCLR": simclr_loss} + + def _simclr_loss(self, x1: torch.Tensor, x2: torch.Tensor): + # simclr training objective from: https://arxiv.org/pdf/2002.05709 + # implementation adapted from: https://github.com/google-research/simclr/blob/master/objective.py + # learnable logit_scale adapted from CLIP: https://arxiv.org/pdf/2103.00020 + + # learnable positive temperature + # TODO may need to do clamping? + logit_scale = self.log_temperature.exp() + + # normalize projections + x1 = F.normalize(x1, p=2, dim=1) # (B, H) + x2 = F.normalize(x2, p=2, dim=1) # (B, H) + B = x1.shape[0] + + # in-batch similarities (should be minimized) + logits_11 = (x1 @ x1.T) * logit_scale + logits_22 = (x2 @ x2.T) * logit_scale + + # mask out autosimilarity diagonal + mask = torch.eye(B, dtype=x1.dtype, device=x1.device) + logits_11 = (1 - mask) * logits_11 + mask * -100 # CE loss ignore_index + logits_22 = (1 - mask) * logits_22 + mask * -100 + + # cross-batch similarities (pairs should be maximized) + logits_12 = (x1 @ x2.T) * logit_scale + logits_21 = (x2 @ x1.T) * logit_scale + + # CE loss over concatenated matrices [A B] + # where A is the cross-batch similiarities whose diagonal should be maximized (given by target) + # and B is the in-batch similarities which should be minimized + labels = torch.arange(B, dtype=torch.long, device=x1.device) + loss_12 = F.cross_entropy( + input=torch.concat((logits_12, logits_11), dim=1), + target=labels, + ) + loss_21 = F.cross_entropy( + input=torch.concat((logits_21, logits_22), dim=1), + target=labels, + ) + + return loss_12 + loss_21 + + @property + def allow_extra_keys(self) -> list[str]: + return [] + + @property + def allow_missing_keys(self) -> list[str]: + return [] diff --git a/protossl/models/_protoecgnet.py b/protossl/models/_protoecgnet.py new file mode 100644 index 0000000..8b06e6b --- /dev/null +++ b/protossl/models/_protoecgnet.py @@ -0,0 +1,452 @@ +from collections import Counter +from dataclasses import dataclass, field + +import torch +import torch.nn as nn +import torch.nn.functional as F +from omegaconf import OmegaConf + +from ..defines import BACKBONE_T, CONV_T, PROT_T, SIM_MAX, STAGE_T +from ._pretrained_utils import PretrainedMixin +from .encoders import PrototypeEncoder + + +class ProtoECGNet(PretrainedMixin, nn.Module): + def __init__( + self, + *, # enforce kwargs + pipeline_stage: STAGE_T, + backbone_type: BACKBONE_T, + conv_type: CONV_T, + prototype_type: PROT_T, + n_prototypes_per_label: int, + label_names: list[str], + label_weights: torch.Tensor, + label_cooccurrence: torch.Tensor | None = None, + pretrained_weights: str | None = None, + partial_len: int | None = None, + partial_overlap: float | None = None, + # ProtoECGNet weights + lam_clst: float = 0.004, + lam_sep: float = 0.0004, + lam_div: float = 250.0, + lam_cntrst: float = 300, + lam_l1: float = 1e-4, + ): + super().__init__() + if ( + pipeline_stage != "learn-prototypes-supervised" + and pipeline_stage != "train-classifier" + ): + raise ValueError( + f"Invalid pipeline_stage for ProtoECGNet: {pipeline_stage}" + ) + if pipeline_stage == "train-classifier" and pretrained_weights is None: + raise ValueError( + f"Stage train-classifier must be used with pretrained weights" + ) + if ( + pipeline_stage == "learn-prototypes-supervised" + and label_cooccurrence is None + ): + raise ValueError( + "Stage learn-prototypes-supervised needs non-None label_cooccurrence" + ) + self.pipeline_stage = pipeline_stage + self.n_prototypes_per_label = n_prototypes_per_label + self.n_labels = len(label_names) + self.label_names = label_names + self.register_buffer("label_weights", label_weights, persistent=False) + self.register_buffer("label_cooccurrence", label_cooccurrence, persistent=False) + self.encoder = PrototypeEncoder( + backbone_type=backbone_type, + n_prototypes=self.n_labels * n_prototypes_per_label, + conv_type=conv_type, + prototype_type=prototype_type, + partial_len=partial_len, + partial_overlap=partial_overlap, + ) + + # mask pos/neg prototype class connections + assign = torch.eye(self.n_labels) # (L, L) + assign = assign.repeat_interleave(n_prototypes_per_label, dim=1) # (L, LP) + off_class_mask = 1.0 - assign # (L, LP) + self.register_buffer("off_class_mask", off_class_mask, persistent=False) + + # binary multilabel classification output + self.cls = nn.Linear( + in_features=self.encoder.emb_dim, + out_features=self.n_labels, + bias=False, + ) + # initialize with 1 for pos connects and -0.5 for neg connects + with torch.no_grad(): + self.cls.weight.copy_(assign - 0.5 * off_class_mask) + + self.lam_clst = lam_clst + self.lam_sep = lam_sep + self.lam_div = lam_div + self.lam_cntrst = lam_cntrst + self.lam_l1 = lam_l1 + + if pretrained_weights is not None: + self.load_pretrained_weights(pretrained_weights) + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple[ + dict[str, torch.Tensor], # losses + dict[str, torch.Tensor], # probs + ]: + sims: torch.Tensor = self.encoder(x) # (B, P) + logits: torch.Tensor = self.cls(sims) # (B, L) + probs = logits.sigmoid() # (B, L) + if self.pipeline_stage == "learn-prototypes-supervised": + losses = self._joint_criterion(sims, logits, y) + elif self.pipeline_stage == "train-classifier": + losses = self._classifier_criterion(sims, logits, y) + else: + raise ValueError( + f"Unknown how to do forward pass for ProtoECGNet during stage {self.pipeline_stage}" + ) + return losses, {label: probs[:, i] for i, label in enumerate(self.label_names)} + + def _classifier_criterion( + self, + sims: torch.Tensor, + logits: torch.Tensor, + y: torch.Tensor, + ) -> dict[str, torch.Tensor]: + # ProtoECGNet: https://arxiv.org/pdf/2504.08713 + # Appendix E, Stage 3 - NOTE: this is just the comparison to the branch-specific classifier + + cls_loss = F.binary_cross_entropy_with_logits( + logits, # (B, L) + y.float(), # (B, L) + pos_weight=self.label_weights, # type: ignore + reduction="mean", + ) + + # masked L1 regularization + penalty = (self.cls.weight.abs() * self.off_class_mask).sum() # type: ignore + + return { + "Classification": cls_loss, + "Penalty": self.lam_l1 * penalty, + } + + def _joint_criterion( + self, + sims: torch.Tensor, + logits: torch.Tensor, + y: torch.Tensor, + ) -> dict[str, torch.Tensor]: + # ProtoECGNet: https://arxiv.org/pdf/2504.08713 + # Appendix E, Stage 1 + + # binary cross entropy loss - no masked l1 penalty + cls_loss = F.binary_cross_entropy_with_logits( + logits, # (B, L) + y.float(), # (B, L) + pos_weight=self.label_weights, # type: ignore + reduction="mean", # NOTE: Eq 3 in ProtoECGNet paper does not show mean reduction over classes but seems like it should be based on their codebase + ) + + # clustering loss + # use repeat_interleave as all prototypes for a given label should be contiguous + pos_mask = y.repeat_interleave(self.n_prototypes_per_label, 1) # (B, P) + has_pos = pos_mask.any(dim=1) # (B,) + neg_mask = 1 - pos_mask + has_neg = neg_mask.any(dim=1) # (B,) + # only consider similarity for prototypes assigned to the sample + # since we take the max of the valid similarities, mask invalid entries + # with similarities less than all other similarities + pos_prot_sims = pos_mask * sims + neg_mask * -SIM_MAX # (B, P) + per_sample_max_pos_sim, _ = pos_prot_sims.max(1) # (B,) + per_sample_max_pos_sim = per_sample_max_pos_sim[has_pos] # (B_p), B_p <= B + clst_loss = -per_sample_max_pos_sim.mean() + + # separation loss + neg_prot_sims = neg_mask * sims + pos_mask * -SIM_MAX # (B, P) + per_sample_max_neg_sim, _ = neg_prot_sims.max(1) # (B,) + per_sample_max_neg_sim = per_sample_max_neg_sim[has_neg] # (B_n), B_n <= B + sep_loss = per_sample_max_neg_sim.mean() + + # orthogonality loss + prots = F.normalize(self.encoder.prototypes, p=2, dim=1) # (P, H) + n_prot = prots.shape[0] + identity = torch.eye(n_prot, device=prots.device) + inter_prot_sims = prots @ prots.T # (P, P) + # squared Frobenius norm (skip the sqrt) of inter-prototype similarities (except self similarity) + # NOTE: division by n_prot^2 not in ProtoECGNet paper but in their codebase + div_loss = ((inter_prot_sims - identity) ** 2).sum() / (n_prot**2) + + # contrastive loss + cooc: torch.Tensor = self.label_cooccurrence # type: ignore - (L, L) + ppl = self.n_prototypes_per_label + # use Kronecker product to expand cooccurrence matrix to prototype assignemnts + cooc_kron = torch.kron(cooc, torch.ones(ppl, ppl, device=cooc.device)) # (P, P) + # NOTE: cooc normalization not in ProtoECGNet paper but in their codebase + pos_cooc = cooc_kron / (cooc_kron.sum() + 1e-6) + neg_cooc = (1 - cooc_kron) / ((1 - cooc_kron).sum() + 1e-6) + pos_weighted_sims = (pos_cooc * inter_prot_sims).sum() + neg_weighted_sims = (neg_cooc * inter_prot_sims).sum() + cntrst_loss = (pos_weighted_sims - neg_weighted_sims) / n_prot**0.5 + + return { + "Classification": cls_loss, + "Clustering": self.lam_clst * clst_loss, + "Separation": self.lam_sep * sep_loss, + "Diversity": self.lam_div * div_loss, + "Contrastive": self.lam_cntrst * cntrst_loss, + } + + @property + def allow_extra_keys(self) -> list[str]: + return [] + + @property + def allow_missing_keys(self) -> list[str]: + return ["cls.weight"] + + +@dataclass +class BranchCfg: + name: str + config: str # path to yaml + pretrained_weights: str | None = None + + # set by post init + backbone_type: BACKBONE_T = field(init=False) + conv_type: CONV_T = field(init=False) + prototype_type: PROT_T = field(init=False) + n_prototypes_per_label: int = field(init=False) + label_subset: list[str] = field(init=False) + partial_len: int | None = field(default=None, init=False) + partial_overlap: float | None = field(default=None, init=False) + + def __post_init__(self): + c = OmegaConf.load(self.config) + OmegaConf.set_struct(c, True) + self.backbone_type = c.model.init_args.backbone_type + self.conv_type = c.model.init_args.conv_type + self.prototype_type = c.model.init_args.prototype_type + self.n_prototypes_per_label = c.model.init_args.n_prototypes_per_label + self.label_subset = c.data.init_args.label_subset + # NOTE: we don't validate correctness, we're just trying to read + # required/optional values from the config yaml here + for opt_k in ["partial_len", "partial_overlap"]: + if opt_k in c.model.init_args: + v = c.model.init_args.get(opt_k) + setattr(self, opt_k, v) + + +BRANCHES_T = list[BranchCfg] + + +def fusion_prototype_assignment( + label_names: list[str], + branches: BRANCHES_T, +) -> torch.Tensor: + """ + return (n_label, n_prototype) binary matrix of label --> prototype assignment, + rows ordered by label_names and columns ordered by prototypes regrouped in label_names order + (i.e. the column order after ProtoECGNetFusion's reverse_mask) + """ + # number of prototypes per label depends on branch that it came from + label_to_ppl = { + label: branch.n_prototypes_per_label + for branch in branches + for label in branch.label_subset + } + total_prototypes = sum(label_to_ppl[label] for label in label_names) + assign = torch.zeros(len(label_names), total_prototypes) # (L, R) + offset = 0 + for row, label in enumerate(label_names): + chunk = label_to_ppl[label] + assign[row, offset : offset + chunk] = 1 + offset += chunk + return assign + + +class ProtoECGNetFusion(PretrainedMixin, nn.Module): + def __init__( + self, + *, # enforce kwargs + pipeline_stage: STAGE_T, + label_names: list[str], # must have same ordering as input y to forward + label_weights: torch.Tensor | None, # only needed for training + branches: BRANCHES_T, + lam_l1: float = 1e-4, + pretrained_weights: str | None = None, + ): + super().__init__() + if pipeline_stage not in { + "train-fusion-classifier", + "compute-fusion-embeddings", + }: + raise ValueError( + f"Invalid pipeline_stage for ProtoECGNetFusion: {pipeline_stage}" + ) + if pipeline_stage == "train-fusion-classifier" and label_weights is None: + raise ValueError( + "label_weights cannot be None for training-fusion-classifier stage" + ) + self.pipeline_stage = pipeline_stage + self.label_names = label_names + total_labels = len(label_names) + self.register_buffer("label_weights", label_weights, persistent=False) + total_encoder_dim, branch_offset = 0, 0 + encoders, label_to_slice = {}, {} + for branch in branches: + if pretrained_weights is None and branch.pretrained_weights is None: + raise ValueError( + "Must specify either fusion checkpoint or per-branch checkpoint" + ) + elif ( + pretrained_weights is not None and branch.pretrained_weights is not None + ): + raise ValueError( + "Cannot provide both fusion checkpoint and per-branch checkpoint" + ) + n_branch_labels = len(branch.label_subset) + branch_ppl = branch.n_prototypes_per_label + encoder = PrototypeEncoder( + backbone_type=branch.backbone_type, + n_prototypes=n_branch_labels * branch_ppl, + conv_type=branch.conv_type, + prototype_type=branch.prototype_type, + partial_len=branch.partial_len, + partial_overlap=branch.partial_overlap, + ) + if branch.pretrained_weights is not None: + # only load encoder weights + sd = torch.load( + branch.pretrained_weights, + weights_only=False, + map_location="cpu", + ) + prefix = "encoder." + if "state_dict" in sd: + sd = sd["state_dict"] + prefix = "model." + prefix + sd = { + k.removeprefix(prefix): v + for k, v in sd.items() + if k.startswith(prefix) + } + encoder.load_state_dict(sd, strict=True) + encoders[branch.name] = encoder + total_encoder_dim += encoder.emb_dim + + # encoder branch --> cls mapping + for branch_idx, label in enumerate(branch.label_subset): + # the column slice in the tensor concatenated from all encoders + label_to_slice[label] = ( + branch_offset + branch_idx * branch_ppl, # start + branch_offset + (branch_idx + 1) * branch_ppl, # end + ) + branch_offset += n_branch_labels * branch_ppl + self.encoders = nn.ModuleDict(encoders) + + # ensure consistency with branch labels + label_count = Counter(label_names) + label_dupes = [x for x, n in label_count.items() if n > 1] + if len(label_dupes) > 0: + raise ValueError(f"label_names contains duplicated labels: {label_dupes}") + branch_count = Counter([l for branch in branches for l in branch.label_subset]) + branch_dupes = [x for x, n in branch_count.items() if n > 1] + if len(branch_dupes) > 0: + raise ValueError(f"branches contain duplicated labels: {branch_dupes}") + missing = set(label_count.keys()) - set(branch_count.keys()) + if len(missing) > 0: + raise ValueError( + f"no branch uses these labels (i.e. these labels are missing): {missing}" + ) + extra = set(branch_count.keys()) - set(label_count.keys()) + if len(extra) > 0: + raise ValueError( + f"some branch uses an unknown label (i.e. these labels are extra): {extra}" + ) + assert branch_offset == total_encoder_dim # quick sanity check + reverse_mask = torch.cat( # (R,) + [ + torch.arange(start, end) + for start, end in (label_to_slice[label] for label in label_names) + ] + ) + self.register_buffer("reverse_mask", reverse_mask, persistent=False) + + # mask pos/neg prototype class connections + # let R be the total number of prototypes across all branches + # the order of the `cls` weights is given by the order of `label_names` + assign = fusion_prototype_assignment(label_names, branches) # (L, R) + assert assign.shape == (total_labels, total_encoder_dim) + off_class_mask = 1.0 - assign # (L, R) + self.register_buffer("off_class_mask", off_class_mask, persistent=False) + + # binary multilabel classification output + # NOTE: deviations from ProtoECGNet codebase: + # * fusion classifier has no bias term here + # * fusion classifier does masked weight init + self.cls = nn.Linear( + in_features=total_encoder_dim, + out_features=total_labels, + bias=False, + ) + # initialize with 1 for pos connects and -0.5 for neg connects + with torch.no_grad(): + self.cls.weight.copy_(assign - 0.5 * off_class_mask) + + self.lam_l1 = lam_l1 + + # this should be a checkpoint for the outer fusion model i.e. all branches and final classifier + if pretrained_weights is not None: + self.load_pretrained_weights(pretrained_weights) + + def compute_prototype_sims(self, x: torch.Tensor) -> tuple[ + torch.Tensor, # (B, n_prototypes) prototype sims, ordered by label_names + torch.Tensor, # (B, n_prototypes) which chunk resulted in each prototype sim + ]: + per_branch_sims, per_branch_chunks = [], [] + for enc in self.encoders.values(): + per_branch_sims.append(enc(x)) + _, chunks = enc.get_last_embs_and_chunks() # type: ignore + per_branch_chunks.append(chunks) + sims: torch.Tensor = torch.concat(per_branch_sims, dim=1) # (B, R) + chunks: torch.Tensor = torch.concat(per_branch_chunks, dim=1) # (B, R) + sims = sims[:, self.reverse_mask] # type: ignore + chunks = chunks[:, self.reverse_mask] # type: ignore + return sims, chunks + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple[ + dict[str, torch.Tensor], # losses + dict[str, torch.Tensor], # probs + ]: + assert self.label_weights is not None, "label_weights is None" + sims, _ = self.compute_prototype_sims(x) # (B, R) + logits: torch.Tensor = self.cls(sims) # (B, L) + probs = logits.sigmoid() # (B, L) + + # ProtoECGNet: https://arxiv.org/pdf/2504.08713 + # Appendix E, Stage 3 + + cls_loss = F.binary_cross_entropy_with_logits( + logits, # (B, L) + y.float(), # (B, L) + pos_weight=self.label_weights, # type: ignore + reduction="mean", + ) + + # masked L1 regularization + penalty = (self.cls.weight.abs() * self.off_class_mask).sum() # type: ignore + + losses = {"Classification": cls_loss, "Penalty": self.lam_l1 * penalty} + labeled_probs = {label: probs[:, i] for i, label in enumerate(self.label_names)} + return losses, labeled_probs + + @property + def allow_extra_keys(self) -> list[str]: + return [] + + @property + def allow_missing_keys(self) -> list[str]: + return [] diff --git a/protossl/models/_prototype_contraster.py b/protossl/models/_prototype_contraster.py index 5a19cec..3559cff 100644 --- a/protossl/models/_prototype_contraster.py +++ b/protossl/models/_prototype_contraster.py @@ -1,3 +1,5 @@ +from typing import get_args + import numpy as np import torch import torch.nn as nn @@ -35,7 +37,7 @@ def __init__( ): super().__init__() - if contrastive_pair_mode not in CONTRASTIVE_T: + if contrastive_pair_mode not in get_args(CONTRASTIVE_T): raise ValueError(f"Unknown contrastive_pair_mode={contrastive_pair_mode}.") if any( [x < 0 for x in [cola_loss_weight, clar_loss_weight, koleo_loss_weight]] diff --git a/protossl/plotting.py b/protossl/plotting.py new file mode 100644 index 0000000..3e20645 --- /dev/null +++ b/protossl/plotting.py @@ -0,0 +1,659 @@ +"""Utilities for plotting ECGs from any of the ECG datasets in this package. + +Waveforms in the datasets are clipped and z-normalized with per-lead (and, for +HEEDB, per-institution) statistics derived over the respective train sets. To +draw a clinical-style printout the waveforms first have to be put back into +physical units, which is what the helpers here take care of, so that the same +plot can be made for a sample from any dataset. + +Typical use:: + + from protossl.datasets import PtbxlECGDataset + from protossl.plotting import plot_ecg + + ds = PtbxlECGDataset(dataset_path=..., split="test", sampling_rate=100) + fig = plot_ecg(dataset=ds, sample_id=3875, chunk_idx=4) # highlight window 4 + fig.savefig("case.png", dpi=300) +""" + +from typing import Type + +import matplotlib.patches as patches +import matplotlib.pyplot as plt +import numpy as np +import torch +from scipy.signal import butter, filtfilt + +from .datasets import ( + BaseTSDataset, + CincECGDataset, + Code15ECGDataset, + EchoNextECGDataset, + HeedbECGDataset, + MimicECGDataset, + PtbxlECGDataset, + ZzuECGDataset, +) +from .datasets._mimic_dataset import standardize_lead_order as _MIMIC_LEAD_PERMUTATION +from .defines import ( + CINC_CLIPPED_MEANS, + CINC_CLIPPED_STDS, + CODE15_CLIPPED_MEANS, + CODE15_CLIPPED_STDS, + HEEDB_EUH_CLIPPED_MEANS, + HEEDB_EUH_CLIPPED_STDS, + HEEDB_MGB_CLIPPED_MEANS, + HEEDB_MGB_CLIPPED_STDS, + MIMIC_CLIPPED_MEANS, + MIMIC_CLIPPED_STDS, + PTBXL_CLIPPED_MEANS, + PTBXL_CLIPPED_STDS, + STANDARD_LEAD_ORDER, + ZZU_CLIPPED_MEANS, + ZZU_CLIPPED_STDS, +) + +# print layout, all datasets emit waveforms in STANDARD_LEAD_ORDER +LEAD_TO_IDX = {name: i for i, name in enumerate(STANDARD_LEAD_ORDER)} +LEAD_LAYOUT = [ + ["I", "aVR", "V1", "V4"], + ["II", "aVL", "V2", "V5"], + ["III", "aVF", "V3", "V6"], +] +RHYTHM_LEADS = ["II"] + +MM_PER_MV = 10 # gain +MM_PER_SEC = 25 # paper speed +ROW_SPACING = 50 # mm between printout rows +SEG_HEIGHT_IN = 1.4 # height of one row of the segment cutout +FULL_SEG_VGAP_REL = ( + 0.1 # vertical gap between printout and cutout, relative to printout +) +PANEL_HGAP_REL = 0.15 # horizontal gap between panels, relative to printout width + +DEFAULT_DURATION = 10.0 # seconds, all ECG datasets emit 10 second recordings + +# per-lead unnormalization stats keyed by dataset class, given in the lead order +# the datasets actually emit. `None` marks data that is not normalized by this +# package (and is therefore already in physical units), a dict marks datasets +# whose samples are normalized with different stats (keyed by `_df["source"]`). +_MIMIC_PERM = np.asarray(_MIMIC_LEAD_PERMUTATION) +_DENORM_STATS: dict[ + Type[BaseTSDataset], + tuple[list[float], list[float]] | dict[str, tuple[list[float], list[float]]] | None, +] = { + PtbxlECGDataset: (PTBXL_CLIPPED_MEANS, PTBXL_CLIPPED_STDS), + CincECGDataset: (CINC_CLIPPED_MEANS, CINC_CLIPPED_STDS), + Code15ECGDataset: (CODE15_CLIPPED_MEANS, CODE15_CLIPPED_STDS), + ZzuECGDataset: (ZZU_CLIPPED_MEANS, ZZU_CLIPPED_STDS), + # MIMIC leads are reordered *after* normalization, so permute the stats the + # exact same way to keep them aligned with the emitted lead order + MimicECGDataset: ( + np.asarray(MIMIC_CLIPPED_MEANS)[_MIMIC_PERM].tolist(), + np.asarray(MIMIC_CLIPPED_STDS)[_MIMIC_PERM].tolist(), + ), + HeedbECGDataset: { + "mgb": (HEEDB_MGB_CLIPPED_MEANS, HEEDB_MGB_CLIPPED_STDS), + "emory": (HEEDB_EUH_CLIPPED_MEANS, HEEDB_EUH_CLIPPED_STDS), + }, + # EchoNext waveforms come prenormalized from the source, nothing to undo + EchoNextECGDataset: None, +} + + +def resolve_index( + dataset: BaseTSDataset, + *, # enforce kwargs + index: int | None = None, + sample_id: int | None = None, +) -> int: + """Resolve a dataset row from either its positional index or its sample ID. + + :param dataset: dataset holding the sample + :type dataset: BaseTSDataset + :param index: positional index into the dataset + :type index: int | None + :param sample_id: sample ID (e.g. PTB-XL `ecg_id`), mutually exclusive with `index` + :type sample_id: int | None + """ + if (index is None) == (sample_id is None): + raise ValueError("Pass exactly one of index or sample_id") + if index is not None: + return int(index) + + # cache the mapping on the dataset, it is a linear scan over all samples + lookup = getattr(dataset, "_sample_id_to_idx", None) + if lookup is None: + lookup = {int(s): i for i, s in enumerate(dataset.sample_ids.tolist())} + dataset._sample_id_to_idx = lookup # type: ignore[attr-defined] + if int(sample_id) not in lookup: # type: ignore[arg-type] + raise KeyError( + f"sample_id {sample_id} not found in this dataset/split " + f"(does it belong to a different split?)" + ) + return lookup[int(sample_id)] # type: ignore[arg-type] + + +def infer_sampling_rate( + dataset: BaseTSDataset, + duration: float = DEFAULT_DURATION, +) -> int: + """Infer the sampling rate from the waveform length, all ECG datasets emit + fixed length recordings of `duration` seconds. + """ + n_timesteps = dataset.waveforms.shape[-1] + sampling_rate = n_timesteps / duration + if sampling_rate != int(sampling_rate): + raise ValueError( + f"Could not infer an integer sampling rate from {n_timesteps} timesteps " + f"over {duration} seconds, pass sampling_rate explicitly" + ) + return int(sampling_rate) + + +def get_denorm_stats( + dataset: BaseTSDataset, + index: int, +) -> tuple[np.ndarray, np.ndarray] | None: + """Per-lead means/stds used to normalize this sample, shaped (L, 1) for + broadcasting over (L, T) waveforms. Returns `None` for datasets whose + waveforms are not normalized by this package. + """ + for cls in type(dataset).__mro__: + if cls in _DENORM_STATS: + stats = _DENORM_STATS[cls] + break + else: + raise ValueError( + f"No unnormalization stats known for {type(dataset).__name__}, " + f"pass means/stds explicitly" + ) + + if stats is None: + return None + if isinstance(stats, dict): + # e.g. HEEDB normalizes per source institution + source = dataset._df.loc[index, "source"] # type: ignore[attr-defined] + if source not in stats: + raise ValueError( + f"{type(dataset).__name__} sample {index} has unknown source " + f"'{source}', expected one of {sorted(stats)}" + ) + stats = stats[source] + + means, stds = stats + return np.asarray(means)[:, None], np.asarray(stds)[:, None] + + +def remove_baseline_wander( + X: np.ndarray, + sampling_rate: int, + cutoff: float = 0.5, + order: int = 1, +) -> np.ndarray: + """Apply a high-pass Butterworth filter to remove baseline wander. + + No need to apply a low-pass filter if using 100 Hz data. + + :param X: waveform of shape (L, T) + :type X: np.ndarray + """ + assert X.ndim == 2 + X = X.astype(np.float64, copy=True) + + b, a = butter(order, cutoff / (sampling_rate / 2), btype="high", analog=False) + for j in range(X.shape[0]): + X[j, :] = filtfilt(b, a, X[j, :]) + + return X + + +def get_ecg_waveform( + *, # enforce kwargs + dataset: BaseTSDataset, + index: int | None = None, + sample_id: int | None = None, + denormalize: bool = True, + remove_baseline: bool = True, + sampling_rate: int | None = None, + duration: float = DEFAULT_DURATION, + means: np.ndarray | list[float] | None = None, + stds: np.ndarray | list[float] | None = None, +) -> np.ndarray: + """Get a single ECG as a (L, T) array in physical units (mV), ready to plot. + + Note that unnormalizing is not exactly the inverse of the dataset transform, + the source waveforms are clipped before normalization. + + :param dataset: any ECG dataset from `protossl.datasets` + :type dataset: BaseTSDataset + :param index: positional index into the dataset, mutually exclusive with `sample_id` + :type index: int | None + :param sample_id: sample ID of the ECG, mutually exclusive with `index` + :type sample_id: int | None + :param denormalize: undo the per-lead z-normalization applied by the dataset + :type denormalize: bool + :param remove_baseline: high-pass filter the waveform to remove baseline wander + :type remove_baseline: bool + :param sampling_rate: sampling rate of the waveform, inferred if not given + :type sampling_rate: int | None + :param duration: recording length in seconds, used to infer the sampling rate + :type duration: float + :param means: per-lead means overriding the dataset's unnormalization stats + :type means: np.ndarray | list[float] | None + :param stds: per-lead stds overriding the dataset's unnormalization stats + :type stds: np.ndarray | list[float] | None + """ + idx = resolve_index(dataset, index=index, sample_id=sample_id) + if sampling_rate is None: + sampling_rate = infer_sampling_rate(dataset, duration) + + waveform = dataset[idx]["waveform"] + if isinstance(waveform, torch.Tensor): + waveform = waveform.detach().cpu().numpy() + signal = np.asarray(waveform, dtype=np.float64) # (L, T) + assert signal.ndim == 2, f"Expected a (L, T) waveform, got {signal.shape}" + + if denormalize: + if (means is None) != (stds is None): + raise ValueError("Pass both means and stds, or neither") + if means is not None and stds is not None: + stats = np.asarray(means)[:, None], np.asarray(stds)[:, None] + else: + stats = get_denorm_stats(dataset, idx) + if stats is not None: + signal = (signal * stats[1]) + stats[0] + + if remove_baseline: + signal = remove_baseline_wander(signal, sampling_rate) + + return signal + + +def chunk_to_seconds( + chunk_idx: int, + *, # enforce kwargs + chunk_len: int, + chunk_overlap: float, + sampling_rate: int, +) -> tuple[float, float]: + """Convert a prototype chunk index into a (start, end) window in seconds. + + Mirrors the windowing in `PrototypeEncoder.forward` for partial prototypes, + where chunks are taken with `unfold(2, partial_len, partial_len * (1 - overlap))`. + + :param chunk_idx: index of the chunk, e.g. from `test_chunks.npy` + :type chunk_idx: int + :param chunk_len: window length in timesteps (the model's `partial_len`) + :type chunk_len: int + :param chunk_overlap: fractional overlap between windows (the model's `partial_overlap`) + :type chunk_overlap: float + :param sampling_rate: sampling rate of the waveform + :type sampling_rate: int + """ + step = int(chunk_len * (1 - chunk_overlap)) + start = (chunk_idx * step) / sampling_rate + return start, start + (chunk_len / sampling_rate) + + +def create_ecg_axes( + n_panels: int = 1, + *, # enforce kwargs + show_segment: bool = True, + duration: float = DEFAULT_DURATION, + lead_layout: list[list[str]] = LEAD_LAYOUT, + rhythm_leads: list[str] = RHYTHM_LEADS, +) -> tuple[plt.Figure, list[tuple[plt.Axes, np.ndarray | None]]]: + """Make a figure with `n_panels` side-by-side printout/cutout panels. + + Each panel is laid out as:: + + [.......... full printout ..........] + [........... vertical gap ..........] + [ seg ][ seg ][ seg ][ seg ] < cutout, one subplot per lead of the layout + + :param n_panels: number of ECGs to show side by side + :type n_panels: int + :param show_segment: include the segment cutout axes below each printout + :type show_segment: bool + :returns: the figure and, per panel, its printout axes and cutout axes grid + """ + if n_panels < 1: + raise ValueError(f"Need at least one panel, got {n_panels}") + + n_seg_rows = len(lead_layout) + n_seg_cols = len(lead_layout[0]) + n_full_rows = n_seg_rows + len(rhythm_leads) + + printout_width_in = (duration * MM_PER_SEC) / 25.4 + printout_height_in = (n_full_rows * ROW_SPACING) / 25.4 + + fig_width_in = printout_width_in * (n_panels + PANEL_HGAP_REL * (n_panels - 1)) + fig_height_in = printout_height_in + height_ratios = [printout_height_in] + if show_segment: + fig_height_in += printout_height_in * FULL_SEG_VGAP_REL + fig_height_in += SEG_HEIGHT_IN * n_seg_rows + height_ratios += [FULL_SEG_VGAP_REL * printout_height_in] + height_ratios += [SEG_HEIGHT_IN] * n_seg_rows + + # per panel: n_seg_cols equal width columns, separated by a gap column + width_ratios: list[float] = [] + for panel in range(n_panels): + if panel > 0: + width_ratios.append(PANEL_HGAP_REL) + width_ratios += [1 / n_seg_cols] * n_seg_cols + + fig = plt.figure(figsize=(fig_width_in, fig_height_in)) + gs = fig.add_gridspec( + nrows=len(height_ratios), + ncols=len(width_ratios), + width_ratios=width_ratios, + height_ratios=height_ratios, + ) + + panels: list[tuple[plt.Axes, np.ndarray | None]] = [] + for panel in range(n_panels): + col_offset = panel * (n_seg_cols + 1) # +1 for the gap column + full_ax = fig.add_subplot(gs[0, col_offset : col_offset + n_seg_cols]) + if not show_segment: + panels.append((full_ax, None)) + continue + + seg_axs: list[list[plt.Axes]] = [] + for i in range(n_seg_rows): + seg_row: list[plt.Axes] = [] + seg_axs.append(seg_row) + for j in range(n_seg_cols): + ax = fig.add_subplot(gs[2 + i, col_offset + j]) + if i != 0 or j != 0: + ax.sharex(seg_axs[0][0]) + seg_row.append(ax) + panels.append((full_ax, np.asarray(seg_axs))) + + return fig, panels + + +def plot_full( + full_signal: np.ndarray, + ax: plt.Axes, + *, # enforce kwargs + sampling_rate: int, + duration: float = DEFAULT_DURATION, + gain: float = MM_PER_MV, + lead_layout: list[list[str]] = LEAD_LAYOUT, + rhythm_leads: list[str] = RHYTHM_LEADS, + lead_to_idx: dict[str, int] = LEAD_TO_IDX, +) -> None: + """Draw a clinical style printout of a (L, T) waveform onto `ax`.""" + n_seg_cols = len(lead_layout[0]) + n_full_rows = len(lead_layout) + len(rhythm_leads) + n_timesteps = full_signal.shape[1] + assert n_timesteps == round(duration * sampling_rate) + full_signal = full_signal.T # (T, L) + + ax.set_xlim(0, duration) + ax.set_ylim(0, n_full_rows * ROW_SPACING) + ax.axis("off") + + # draw paper grid, 1 small box = 0.04 sec x 1 mm, 1 large box = 5 small boxes + for x in np.arange(0, duration + 0.04, 0.04): + ax.axvline(x, color="pink", linewidth=0.5, zorder=0) + for y in np.arange(0, n_full_rows * ROW_SPACING + 0.1, 1): + is_large = (y * 0.1) % 0.5 == 0 + ax.axhline( + y, + color="pink" if not is_large else "red", + linewidth=0.5 if not is_large else 1.0, + zorder=0, + ) + for x in np.arange(0, duration + 0.2, 0.2): + ax.axvline(x, color="red", linewidth=1.0, zorder=0) + + label_fontsize = 12 + row_baseline = n_full_rows * ROW_SPACING - ROW_SPACING / 2 + col_duration = duration / n_seg_cols + col_timesteps = n_timesteps // n_seg_cols + + for row_idx, lead_row in enumerate(lead_layout): + for col_idx, lead in enumerate(lead_row): + lead_idx = lead_to_idx[lead] + start = col_idx * col_timesteps + end = start + col_timesteps + signal = full_signal[start:end, lead_idx] * gain + t = np.linspace( + col_idx * col_duration, + (col_idx + 1) * col_duration, + col_timesteps, + ) + v_offset = row_baseline - row_idx * ROW_SPACING + ax.plot(t, signal + v_offset, color="black", linewidth=1.0) + ax.text( + t[0] + 0.1, + v_offset + 16, + lead, + fontsize=label_fontsize, + fontweight="bold", + ) + + # rhythm leads: the full recording + t = np.linspace(0, duration, n_timesteps) + for i, lead in enumerate(rhythm_leads): + v_offset = row_baseline - (len(lead_layout) + i) * ROW_SPACING + signal = full_signal[:, lead_to_idx[lead]] * gain + ax.plot(t, signal + v_offset, color="black", linewidth=1.0) + ax.text( + 0.1, + v_offset + 16, + lead, + fontsize=label_fontsize, + fontweight="bold", + ) + + +def plot_segment( + segment_signal: np.ndarray, + axs: np.ndarray, + *, # enforce kwargs + sampling_rate: int, + lead_layout: list[list[str]] = LEAD_LAYOUT, + lead_to_idx: dict[str, int] = LEAD_TO_IDX, + label_leads: bool = True, +) -> None: + """Draw a cutout of a (L, T) segment, one subplot per lead of the layout. + + Pass `label_leads=False` when the caller labels the leads itself, e.g. as + column headers of a grid of cutouts. + """ + n_rows = len(lead_layout) + n_cols = len(lead_layout[0]) + assert axs.shape == (n_rows, n_cols) # to match layout + t_steps = segment_signal.shape[1] + t_segment = np.linspace(0, t_steps / sampling_rate, t_steps) + + for row in range(n_rows): + for col in range(n_cols): + lead = lead_layout[row][col] + i = lead_to_idx[lead] + ax: plt.Axes = axs[row, col] + ax.plot(t_segment, segment_signal[i], color="black", linewidth=1.0) + ax.set_xlim(0, t_segment[-1]) + ax.set_ylim( + np.min(segment_signal[i]) - 0.5, np.max(segment_signal[i]) + 0.5 + ) + + # draw lead label to the left of the plot + if label_leads: + ax.text( + x=-0.1, + y=0.8, + s=lead, + fontsize=10, + fontweight="bold", + transform=ax.transAxes, + ha="right", + va="bottom", + ) + + # ECG grid + for x in np.arange(t_segment[0], t_segment[-1] + 0.04, 0.04): + ax.axvline(x, color="pink", linewidth=0.3, alpha=0.5) + for x in np.arange(t_segment[0], t_segment[-1], 0.2): + ax.axvline(x, color="red", linewidth=0.5, alpha=0.7) + ax.axhline(0, color="red", linewidth=0.5, alpha=0.5) + + ax.tick_params(left=False, labelleft=False, bottom=False, labelbottom=False) + + +def plot_ecg( + *, # enforce kwargs + dataset: BaseTSDataset | None = None, + index: int | None = None, + sample_id: int | None = None, + signal: np.ndarray | None = None, + highlight: tuple[float, float] | None = None, + chunk_idx: int | None = None, + chunk_len: int | None = None, + chunk_overlap: float = 0.5, + show_segment: bool = True, + title: str | None = None, + title_kwargs: dict | None = None, + full_ax: plt.Axes | None = None, + seg_axs: np.ndarray | None = None, + duration: float = DEFAULT_DURATION, + sampling_rate: int | None = None, + denormalize: bool = True, + remove_baseline: bool = True, + gain: float = MM_PER_MV, + lead_layout: list[list[str]] = LEAD_LAYOUT, + rhythm_leads: list[str] = RHYTHM_LEADS, + means: np.ndarray | list[float] | None = None, + stds: np.ndarray | list[float] | None = None, +) -> plt.Figure: + """Plot a single ECG from any ECG dataset as a clinical style printout, + optionally highlighting a window and calling it out as a cutout below. + + :param dataset: any ECG dataset from `protossl.datasets`, optional only if + both `signal` and `sampling_rate` are given + :type dataset: BaseTSDataset | None + :param index: positional index into the dataset, mutually exclusive with `sample_id` + :type index: int | None + :param sample_id: sample ID of the ECG, mutually exclusive with `index` + :type sample_id: int | None + :param signal: preloaded (L, T) waveform in mV, skips loading from the dataset + :type signal: np.ndarray | None + :param highlight: window to highlight as (start, end) in seconds + :type highlight: tuple[float, float] | None + :param chunk_idx: prototype chunk to highlight, mutually exclusive with `highlight` + :type chunk_idx: int | None + :param chunk_len: window length in timesteps (the model's `partial_len`), + defaults to one second + :type chunk_len: int | None + :param chunk_overlap: fractional overlap between windows (the model's `partial_overlap`) + :type chunk_overlap: float + :param show_segment: draw the cutout of the highlighted window below the printout + :type show_segment: bool + :param full_ax: existing axes to draw the printout on, e.g. from `create_ecg_axes`, + a new single panel figure is made if not given + :type full_ax: plt.Axes | None + :param seg_axs: existing cutout axes matching the lead layout, required when + drawing the cutout onto `full_ax` + :type seg_axs: np.ndarray | None + :returns: the figure that was drawn on + """ + if dataset is None and (signal is None or sampling_rate is None): + raise ValueError("Pass a dataset, or both signal and sampling_rate") + if sampling_rate is None: + assert dataset is not None + sampling_rate = infer_sampling_rate(dataset, duration) + if chunk_idx is not None: + if highlight is not None: + raise ValueError("Pass exactly one of highlight or chunk_idx") + highlight = chunk_to_seconds( + chunk_idx, + chunk_len=chunk_len if chunk_len is not None else sampling_rate, + chunk_overlap=chunk_overlap, + sampling_rate=sampling_rate, + ) + show_segment = show_segment and highlight is not None + + if signal is None: + assert dataset is not None + signal = get_ecg_waveform( + dataset=dataset, + index=resolve_index(dataset, index=index, sample_id=sample_id), + denormalize=denormalize, + remove_baseline=remove_baseline, + sampling_rate=sampling_rate, + duration=duration, + means=means, + stds=stds, + ) + + made_fig = full_ax is None + if full_ax is None: + fig, panels = create_ecg_axes( + 1, + show_segment=show_segment, + duration=duration, + lead_layout=lead_layout, + rhythm_leads=rhythm_leads, + ) + full_ax, seg_axs = panels[0] + else: + fig = full_ax.get_figure() + if show_segment and seg_axs is None: + raise ValueError("Pass seg_axs alongside full_ax to draw the cutout") + + plot_full( + signal, + full_ax, + sampling_rate=sampling_rate, + duration=duration, + gain=gain, + lead_layout=lead_layout, + rhythm_leads=rhythm_leads, + ) + + if highlight is not None: + highlight_start, highlight_end = highlight + full_ax.axvspan( + highlight_start, highlight_end, color="blue", alpha=0.2, zorder=1 + ) + + if show_segment: + assert highlight is not None and seg_axs is not None + ymin, _ = full_ax.get_ylim() + raw_start = max(int(highlight[0] * sampling_rate), 0) + raw_end = min(int(highlight[1] * sampling_rate), signal.shape[1]) + plot_segment( + signal[:, raw_start:raw_end], + seg_axs, + sampling_rate=sampling_rate, + lead_layout=lead_layout, + ) + + # callout lines from the highlighted window to the cutout + for x_full, x_seg, seg_ax in [ + (highlight[0], 0, seg_axs[0, 0]), # bottom left of window to cutout + (highlight[1], 1, seg_axs[0, -1]), # bottom right of window to cutout + ]: + fig.add_artist( + patches.ConnectionPatch( + xyA=(x_full, ymin), + coordsA="data", + xyB=(x_seg, 1), # top corner of the cutout + coordsB="axes fraction", + axesA=full_ax, + axesB=seg_ax, + color="black", + linewidth=1.5, + ) + ) + + if title is not None: + full_ax.set_title(title, **(title_kwargs or {})) + + if made_fig: + fig.tight_layout() + return fig diff --git a/protossl/protoecgnet_trainer.py b/protossl/protoecgnet_trainer.py new file mode 100644 index 0000000..675989c --- /dev/null +++ b/protossl/protoecgnet_trainer.py @@ -0,0 +1,777 @@ +import os + +import numpy as np +import pandas as pd +import torch +from lightning import LightningDataModule, LightningModule +from lightning.pytorch.callbacks import BasePredictionWriter +from lightning.pytorch.cli import LightningCLI +from lightning.pytorch.strategies import DDPStrategy, SingleDeviceStrategy +from torch.utils.data import DataLoader +from torchmetrics import MetricCollection +from torchmetrics.classification import ( + BinaryAUROC, + BinaryAveragePrecision, + MultilabelAUROC, + MultilabelAveragePrecision, +) + +from .datasets import infer_dataset_class_from_path, validate_label_subset +from .defines import BACKBONE_T, CONV_T, PROT_T, SIM_MAX, STAGE_T +from .lightning_utils import check_final_link +from .models import PrototypeProjector +from .models._protoecgnet import BRANCHES_T, ProtoECGNet, ProtoECGNetFusion + +torch.set_float32_matmul_precision("medium") + +PROTOECGNET_PIPELINE_STAGES = [ + "learn-prototypes-supervised", + "project-prototypes-supervised", + "compute-embeddings", + "train-classifier", + "train-fusion-classifier", + "compute-fusion-embeddings", +] + +ProtoECGNetTrainerError = ValueError( + f"This specialized trainer should only be used for ProtoECGNet-style modeling, " + f"restricted to the following pipeline_stages: {PROTOECGNET_PIPELINE_STAGES}" +) + + +class LitData(LightningDataModule): + def __init__( + self, + dataset_path: str, + pipeline_stage: STAGE_T, + batch_size: int, + num_workers: int, + sampling_rate: int = 100, + prefetch_factor: int | None = None, + label_subset: list[str] | None = None, + data_kwargs: dict = dict(), + ): + super().__init__() + self.save_hyperparameters() + + if pipeline_stage not in PROTOECGNET_PIPELINE_STAGES: + raise ProtoECGNetTrainerError + + # label names is linked to LitModel via LitCLI + self.ds_cls, self.label_names, self.is_audio = infer_dataset_class_from_path( + dataset_path + ) + if label_subset is not None: + if self.label_names is None: + raise ValueError( + f"Can only specify label_subset for a dataset with defined labels" + ) + validate_label_subset(label_subset, self.label_names) + self.label_names = label_subset + + if ( + pipeline_stage == "learn-prototypes-supervised" + or pipeline_stage == "train-classifier" + or pipeline_stage == "train-fusion-classifier" + ): + # need label weights and cooccurrence matrix, not best practice + # to instantiate dataset outside of `setup` but we know we'll need + # the train_ds in this pipeline_stage anyways + self.train_ds = self.ds_cls( + dataset_path=dataset_path, + split="train", + sampling_rate=sampling_rate, + label_subset=label_subset, + **data_kwargs, + ) + self.label_weights = self.train_ds.get_label_weights() + self.label_cooccurrence = self.train_ds.get_label_cooccurrence() + else: + # set to defaults so LitCLI can link these to LitModel args + self.label_weights = None + self.label_cooccurrence = None + + def setup(self, stage: str): + dataset_path: str = self.hparams.dataset_path # type: ignore + sampling_rate: int = self.hparams.sampling_rate # type: ignore + label_subset: list[str] | None = self.hparams.label_subset # type: ignore + data_kwargs: dict = self.hparams.data_kwargs # type: ignore + + if stage == "fit": + if not hasattr(self, "train_ds"): + self.train_ds = self.ds_cls( + dataset_path=dataset_path, + split="train", + sampling_rate=sampling_rate, + label_subset=label_subset, + **data_kwargs, + ) + else: + assert self.train_ds is not None, "Not sure how train_ds is None" + + if stage in ["fit", "validate"]: + self.val_ds = self.ds_cls( + dataset_path=dataset_path, + split="val", + sampling_rate=sampling_rate, + label_subset=label_subset, + **data_kwargs, + ) + + if stage in ["test", "predict"]: + self.test_ds = self.ds_cls( + dataset_path=dataset_path, + split="test", + sampling_rate=sampling_rate, + label_subset=label_subset, + **data_kwargs, + ) + + def train_dataloader(self): + pipeline_stage: STAGE_T = self.hparams.pipeline_stage # type: ignore + + return DataLoader( + self.train_ds, + shuffle=( + # no shuffle if projecting/embedding + pipeline_stage + not in { + "project-prototypes-supervised", + "compute-embeddings", + "compute-fusion-embeddings", + } + ), + pin_memory=True, + drop_last=False, + batch_size=self.hparams.batch_size, # type: ignore + num_workers=self.hparams.num_workers, # type: ignore + prefetch_factor=self.hparams.prefetch_factor, # type: ignore + ) + + def val_dataloader(self): + return DataLoader( + self.val_ds, + shuffle=False, + pin_memory=True, + drop_last=False, + batch_size=self.hparams.batch_size, # type: ignore + num_workers=self.hparams.num_workers, # type: ignore + prefetch_factor=self.hparams.prefetch_factor, # type: ignore + ) + + def test_dataloader(self): + return DataLoader( + self.test_ds, + shuffle=False, + pin_memory=True, + drop_last=False, + batch_size=self.hparams.batch_size, # type: ignore + num_workers=self.hparams.num_workers, # type: ignore + prefetch_factor=self.hparams.prefetch_factor, # type: ignore + ) + + def predict_dataloader(self): + return self.test_dataloader() + + +class LitModel(LightningModule): + def __init__( + self, + pipeline_stage: STAGE_T, + label_names: list[str], + backbone_type: BACKBONE_T | None = None, + conv_type: CONV_T | None = None, + prototype_type: PROT_T | None = None, + n_prototypes_per_label: int | None = None, + label_weights: torch.Tensor | None = None, + label_cooccurrence: torch.Tensor | None = None, + pretrained_weights: str | None = None, + partial_len: int | None = None, + partial_overlap: float | None = None, + lam_clst: float = 0.004, + lam_sep: float = 0.0004, + lam_div: float = 250.0, + lam_cntrst: float = 300, + lam_l1: float = 1e-4, + branches: BRANCHES_T | None = None, + ): + super().__init__() + self.lr = None + self.save_hyperparameters() + n_labels = len(label_names) + if n_labels == 1: + metrics = MetricCollection( + { + "ap": BinaryAveragePrecision(), + "auroc": BinaryAUROC(), + } + ) + else: + metrics = MetricCollection( + { + "ap": MultilabelAveragePrecision( + num_labels=n_labels, average="macro" + ), + "auroc": MultilabelAUROC(num_labels=n_labels, average="macro"), + } + ) + self.val_metrics = metrics.clone(prefix="val_") + self.test_metrics = metrics.clone(prefix="test_") + + if pipeline_stage not in PROTOECGNET_PIPELINE_STAGES: + raise ProtoECGNetTrainerError + + n_prototypes = None + if pipeline_stage == "learn-prototypes-supervised": + assert label_weights is not None, "label_weights is None" + assert label_cooccurrence is not None, "label_cooccurrence is None" + assert backbone_type is not None, "backbone_type is None" + assert conv_type is not None, "conv_type is None" + assert prototype_type is not None, "prototype_type is None" + assert n_prototypes_per_label is not None, "n_prototypes_per_label is None" + # don't enforce default at litmodel, let model handle defaults + self.model = ProtoECGNet( + pipeline_stage=pipeline_stage, + backbone_type=backbone_type, + conv_type=conv_type, + prototype_type=prototype_type, + n_prototypes_per_label=n_prototypes_per_label, + label_names=label_names, + label_weights=label_weights, + label_cooccurrence=label_cooccurrence, + pretrained_weights=pretrained_weights, + partial_len=partial_len, + partial_overlap=partial_overlap, + lam_clst=lam_clst, + lam_sep=lam_sep, + lam_div=lam_div, + lam_cntrst=lam_cntrst, + lam_l1=lam_l1, + ) + elif ( + pipeline_stage == "project-prototypes-supervised" + or pipeline_stage == "compute-embeddings" + ): + self._batch_counter = 0 + assert pretrained_weights is not None, "pretrained_weights is None" + assert backbone_type is not None, "backbone_type is None" + assert conv_type is not None, "conv_type is None" + assert prototype_type is not None, "prototype_type is None" + assert n_prototypes_per_label is not None, "n_prototypes_per_label is None" + n_prototypes = n_prototypes_per_label * len(label_names) + self.model = PrototypeProjector( + backbone_type=backbone_type, + conv_type=conv_type, + prototype_type=prototype_type, + n_prototypes=n_prototypes, + pretrained_weights=pretrained_weights, + partial_len=partial_len, + partial_overlap=partial_overlap, + ) + elif pipeline_stage == "train-classifier": + assert label_weights is not None, "label_weights is None" + assert pretrained_weights is not None, "pretrained_weights is None" + assert backbone_type is not None, "backbone_type is None" + assert conv_type is not None, "conv_type is None" + assert prototype_type is not None, "prototype_type is None" + assert n_prototypes_per_label is not None, "n_prototypes_per_label is None" + # don't enforce default at litmodel, let model handle defaults + self.model = ProtoECGNet( + pipeline_stage=pipeline_stage, + backbone_type=backbone_type, + conv_type=conv_type, + prototype_type=prototype_type, + n_prototypes_per_label=n_prototypes_per_label, + label_names=label_names, + label_weights=label_weights, + pretrained_weights=pretrained_weights, + partial_len=partial_len, + partial_overlap=partial_overlap, + lam_clst=lam_clst, + lam_sep=lam_sep, + lam_div=lam_div, + lam_cntrst=lam_cntrst, + lam_l1=lam_l1, + ) + for param in self.model.encoder.parameters(): + param.requires_grad = False + elif pipeline_stage == "train-fusion-classifier": + assert label_weights is not None, "label_weights is None" + assert branches is not None, "branches is None" + # don't enforce default at litmodel, let model handle defaults + self.model = ProtoECGNetFusion( + pipeline_stage=pipeline_stage, + label_names=label_names, + label_weights=label_weights, + branches=branches, + pretrained_weights=pretrained_weights, + lam_l1=lam_l1, + ) + for param in self.model.encoders.parameters(): + param.requires_grad = False + elif pipeline_stage == "compute-fusion-embeddings": + assert branches is not None, "branches is None" + # pretrained_weights (fusion ckpt) or per-branch weights validated by model + self.model = ProtoECGNetFusion( + pipeline_stage=pipeline_stage, + label_names=label_names, + label_weights=label_weights, + branches=branches, + pretrained_weights=pretrained_weights, + lam_l1=lam_l1, + ) + for param in self.model.parameters(): + param.requires_grad = False + else: + raise ValueError(f"Unknown pipeline_stage {pipeline_stage}") + + if pipeline_stage == "project-prototypes-supervised": + assert n_prototypes is not None + # placeholder tensor values for the prediction writer to access, mostly to help with type checking + self.prototype_sims = torch.ones((n_prototypes,)) * -SIM_MAX + self.prototype_embs = torch.empty( + (n_prototypes, self.model.encoder.backbone.emb_dim) # type: ignore + ) + # each row is source_id, sample_id, chunk_idx + self.prototype_ids = torch.empty((n_prototypes, 3), dtype=torch.long) + + def _common_step( + self, + *, # enforce kwargs + batch: dict[str, torch.Tensor], + stage: str, + log: bool = True, + sync_dist: bool = False, + ) -> tuple[ + torch.Tensor | None, # scalar (composite) loss value + dict[str, torch.Tensor] # 1D arrays of predicted probabilities + | tuple[torch.Tensor, torch.Tensor] # embeddings, metadata + | None, + ]: + batch_size = batch["source_id"].shape[0] + pipeline_stage: STAGE_T = self.hparams.pipeline_stage # type: ignore + + if ( + pipeline_stage == "learn-prototypes-supervised" + or pipeline_stage == "train-classifier" + or pipeline_stage == "train-fusion-classifier" + ): + loss_terms, preds = self.model(batch["waveform"], batch["label"]) + loss = self._log_and_composite_losses( + stage=stage, + losses=loss_terms, + batch_size=batch_size, + log=log, + sync_dist=sync_dist, + ) + probs = torch.stack([preds[l] for l in self.hparams.label_names], dim=1) # type: ignore + if stage == "val": + self.val_metrics.update(probs, batch["label"]) + elif stage == "test": + self.test_metrics.update(probs, batch["label"]) + elif pipeline_stage == "project-prototypes-supervised": + self._batch_counter += 1 + assert isinstance(self.model, PrototypeProjector) + if stage != "predict": + raise ValueError( + f"Cannot use pipeline_stage=project-prototypes with non-predict stage (got stage={stage}).\n" + f"We hijack lightning.Trainer.predict with the train set for prototype projection (model in eval mode and write projection metadata)." + ) + loss, preds = None, None + # waveform/label keys + sims = self.model(batch["waveform"]) # (B, n_prototypes) + if pipeline_stage == "project-prototypes-supervised": + # mask out prototypes for which the sample cannot belong to + _, L = batch["label"].shape # (B, L) + _, P = sims.shape + ppl = P // L # n_prototypes_per_label + pos_mask = batch["label"].repeat_interleave(ppl, 1) # (B, P) + sims = pos_mask * sims + (1 - pos_mask) * -SIM_MAX # (B, P) + ( + embs, # (B, chunks, d_emb) + chunks, # (B, n_prototypes) - which chunks resulted in the prototype sims? + ) = self.model.get_last_embs_and_chunks() + + # move things off gpu + sims, embs, chunks = ( + sims.detach().cpu(), + embs.detach().cpu(), + chunks.detach().cpu(), + ) + source_ids, sample_ids = batch["source_id"].cpu(), batch["sample_id"].cpu() + + # find max sim + for prot_idx, curr_sim in enumerate(self.prototype_sims): + prot_sims = sims[:, prot_idx] + candidates = (prot_sims > curr_sim).argwhere().squeeze(1) + if candidates.shape[0] == 0: + # none in batch are more similar to any of the prototypes + continue + # get the candidate in batch that is the most similar + batch_idx = candidates[prot_sims[candidates].argmax()] + chunk_idx = chunks[batch_idx, prot_idx] + + self.prototype_embs[prot_idx] = embs[batch_idx, chunk_idx] + self.prototype_sims[prot_idx] = prot_sims[batch_idx] + self.prototype_ids[prot_idx, 0] = source_ids[batch_idx] + self.prototype_ids[prot_idx, 1] = sample_ids[batch_idx] + self.prototype_ids[prot_idx, 2] = chunk_idx + elif pipeline_stage == "compute-embeddings": + assert isinstance(self.model, PrototypeProjector) + if stage != "predict": + raise ValueError( + f"Cannot use pipeline_stage=compute-embeddings with non-predict stage (got stage={stage}).\n" + ) + loss = None + sims = self.model(batch["waveform"]) # (B, n_prototypes) + # which chunks resulted in the prototype sims? + _, chunks = self.model.get_last_embs_and_chunks() # (B, n_prototypes) + + preds = sims, chunks + elif pipeline_stage == "compute-fusion-embeddings": + assert isinstance(self.model, ProtoECGNetFusion) + if stage != "predict": + raise ValueError( + f"Cannot use pipeline_stage=compute-fusion-embeddings with non-predict stage (got stage={stage}).\n" + ) + loss = None + # (B, R) sims and chunks, concatenated across branches in label_names order + sims, chunks = self.model.compute_prototype_sims(batch["waveform"]) + preds = sims, chunks + else: + raise ValueError( + f"Unknown forward step for pipeline_stage {pipeline_stage}" + ) + return loss, preds + + def _log_and_composite_losses( + self, + *, # enforce kwargs + stage: str, + losses: dict[str, torch.Tensor], + batch_size: int, + log: bool = True, + sync_dist: bool = False, + ) -> torch.Tensor: + for k, v in losses.items(): + if log: + self.log( + f"{stage}_{k}_loss", v, batch_size=batch_size, sync_dist=sync_dist + ) + loss = torch.stack(list(losses.values())).sum() # NOTE: sum, not mean + if log: + self.log(f"{stage}_loss", loss, batch_size=batch_size, sync_dist=sync_dist) + return loss + + def training_step(self, batch, batch_idx): + loss, _ = self._common_step(batch=batch, stage="train") + return loss + + def validation_step(self, batch, batch_idx): + loss, _ = self._common_step(batch=batch, stage="val", sync_dist=True) + return loss + + def on_validation_epoch_end(self): + metrics = self.val_metrics.compute() + self.log_dict(metrics) + self.val_metrics.reset() + + def test_step(self, batch, batch_idx): + loss, _ = self._common_step(batch=batch, stage="test", sync_dist=True) + return loss + + def on_test_epoch_end(self): + metrics = self.test_metrics.compute() + self.log_dict(metrics) + self.test_metrics.reset() + + def predict_step(self, batch, batch_idx): + _, preds = self._common_step(batch=batch, stage="predict", log=False) + return preds + + def on_predict_batch_end(self, outputs, batch, batch_idx, dataloader_idx=0): + pipeline_stage: STAGE_T = self.hparams.pipeline_stage # type: ignore + if pipeline_stage == "project-prototypes-supervised": + assert isinstance(self.model, PrototypeProjector) + + # sync projected prototypes across distributed ranks + # NOTE: for speed, we only run this once (in the last batch) + # it would be cleaner to put this in on_predict_end, however the + # prediction writer's on_predict_end hook fires before the model's + # and we shouldn't really rely on enforcing peer callbacks' order + if ( + len(self.trainer.predict_dataloaders) == batch_idx + 1 # type: ignore + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ): + world_size = torch.distributed.get_world_size() + # fmt: off + gathered_embs = [None for _ in range(world_size)] + gathered_sims = [None for _ in range(world_size)] + gathered_ids = [None for _ in range(world_size)] + torch.distributed.all_gather_object(gathered_embs, self.prototype_embs) + torch.distributed.all_gather_object(gathered_sims, self.prototype_sims) + torch.distributed.all_gather_object(gathered_ids, self.prototype_ids) + gathered_embs = torch.stack(gathered_embs) # type: ignore + gathered_sims = torch.stack(gathered_sims) # type: ignore + gathered_ids = torch.stack(gathered_ids) # type: ignore + torch.distributed.barrier() + # fmt: on + + _, idxs = gathered_sims.max(dim=0) # (P,) + for p in range(len(idxs)): + idx = idxs[p] # for proto p, rank with max sim + self.prototype_embs[p] = gathered_embs[idx][p] # type: ignore + self.prototype_sims[p] = gathered_sims[idx][p] # type: ignore + self.prototype_ids[p] = gathered_ids[idx][p] # type: ignore + + def on_predict_end(self): + # if we've hijacked predict to do prototype projection, need to set + # the prototypes to the projected samples - metadata is saved by + # prediction writer utilities and checkpoint is saved in main function + pipeline_stage: STAGE_T = self.hparams.pipeline_stage # type: ignore + if pipeline_stage == "project-prototypes-supervised": + assert isinstance(self.model, PrototypeProjector) + + # save projected prototypes to model parameter + with torch.no_grad(): + self.model.encoder.prototypes.copy_(self.prototype_embs) + + +class PredictionWriter(BasePredictionWriter): + def __init__(self): + super().__init__("epoch") + + def write_on_epoch_end( + self, + trainer, + pl_module, + predictions, + batch_indices, + ): + # hard to set pipeline_stage on prediction_writer with CLI link_arguments but we can borrow from lit module + pipeline_stage: STAGE_T = pl_module.hparams.pipeline_stage # type: ignore + + if pipeline_stage == "project-prototypes-supervised": + assert hasattr(pl_module, "prototype_ids") + # predictions are all None (see LitModel above) + # pl_module has prototype projection metadata to save + meta = pd.DataFrame.from_dict( + { + "prototype_id": np.arange(len(pl_module.prototype_sims)), # type: ignore + "source_id": pl_module.prototype_ids[:, 0].tolist(), # type: ignore + "sample_id": pl_module.prototype_ids[:, 1].tolist(), # type: ignore + "chunk_idx": pl_module.prototype_ids[:, 2].tolist(), # type: ignore + "emb_sim": pl_module.prototype_sims.tolist(), # type: ignore + }, + orient="columns", + ) + # after gathering data to save, only save on rank 0 + if trainer.global_rank == 0: + meta.to_csv( + os.path.join(trainer.log_dir, "projection_metadata.csv"), # type: ignore + index=False, + ) + to_save = None # type: ignore + save_name = None + elif ( + pipeline_stage == "train-classifier" + or pipeline_stage == "train-fusion-classifier" + ): + # TODO batch order when using distributed + assert isinstance(predictions[0], dict) # classifiction probabilities + if "Multiclass" in predictions[0]: + to_save = np.concatenate([p["Multiclass"] for p in predictions]) + else: + target_names: list[str] = pl_module.hparams.label_names # type: ignore + probs = {k: [] for k in target_names} + for batch_probs in predictions: + for k in target_names: + batch_label_prob: torch.Tensor = batch_probs[k] + probs[k].append(batch_label_prob.numpy()) + to_save = np.stack([np.concatenate(v) for v in probs.values()]).T + if ( + hasattr(pl_module, "prediction_split") + and pl_module.prediction_split is not None + ): + save_name = f"{pl_module.prediction_split}_probs.npy" + else: + save_name = "probs.npy" + elif ( + pipeline_stage == "compute-embeddings" + or pipeline_stage == "compute-fusion-embeddings" + ): + assert ( + hasattr(pl_module, "prediction_split") + and pl_module.prediction_split is not None + ) + split = pl_module.prediction_split + assert isinstance(predictions[0], tuple) + assert isinstance(predictions[0][0], torch.Tensor) # embeddings + assert isinstance(predictions[0][1], torch.Tensor) # chunk metadata + embeddings = [p[0] for p in predictions] + to_save = torch.concatenate(embeddings).numpy() # type: ignore + save_name = f"{split}_embeds.npy" + + # TODO: gather/save chunk metadata with DDP + chunk_metadata = [p[1] for p in predictions] + chunk_metadata = torch.concatenate(chunk_metadata).numpy() + np.save(os.path.join(trainer.log_dir, f"{split}_chunks.npy"), chunk_metadata) # type: ignore + else: + raise ValueError( + f"Unknown how to handle predictions for pipeline_stage={pipeline_stage}" + ) + + # nothing to save or save handled separately + if to_save is None: + return + + # NOTE: currently only single device and DDP strategies have been validated + # saving (potentially distributed) predictions + + # batch_indices is a singleton list of a list of lists + # e.g. [[[0, 2], [1, 3]]], where there are 4 samples and 2 batches of 2 samples each + assert len(batch_indices) == 1 # do other strategies return different idxs? + flat_idxs = np.asarray([x for xs in batch_indices[0] for x in xs]) + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + gathered_to_save = [None] * world_size + gathered_flat_idxs = [None] * world_size + torch.distributed.all_gather_object(gathered_to_save, to_save) + torch.distributed.all_gather_object(gathered_flat_idxs, flat_idxs) + torch.distributed.barrier() + to_save: np.ndarray = np.concatenate(gathered_to_save) # type: ignore + flat_idxs: np.ndarray = np.concatenate(gathered_flat_idxs) # type: ignore + sort_mask = flat_idxs.argsort() + to_save = to_save[sort_mask] + + # after gathering data to save, only save on rank 0 + if trainer.global_rank != 0: + return + log_dir: str = trainer.log_dir # type: ignore + np.save(os.path.join(log_dir, save_name), to_save) # type: ignore + + +class LitCLI(LightningCLI): + def add_arguments_to_parser(self, parser): + parser.add_argument( + "--pipeline-stage", + choices=PROTOECGNET_PIPELINE_STAGES, + required=True, + ) + parser.add_argument("--resume-from-checkpoint") + parser.link_arguments( + "resume_from_checkpoint", "trainer.logger.init_args.resume_from_checkpoint" + ) + parser.link_arguments("pipeline_stage", "data.init_args.pipeline_stage") + parser.link_arguments("pipeline_stage", "model.init_args.pipeline_stage") + parser.link_arguments( + "pipeline_stage", + "trainer.logger.init_args.pipeline_stage", + ) + parser.link_arguments( + "data.label_names", + "model.init_args.label_names", + apply_on="instantiate", + ) + parser.link_arguments( + "data.label_weights", + "model.init_args.label_weights", + apply_on="instantiate", + ) + parser.link_arguments( + "data.label_cooccurrence", + "model.init_args.label_cooccurrence", + apply_on="instantiate", + ) + + +def run(): + cli = LitCLI( + run=False, + save_config_kwargs={"overwrite": True}, # for resuming checkpoint in same dir + ) + _seed = os.environ.get("PL_GLOBAL_SEED", None) + if _seed is not None: + if int(_seed) < 42: + # this is a silly limitation tied to our initial results + # where ILP prototype assignment had random seed = 0 while the rest + # of the initial results used random seed 42 elsewhere. to ensure + # reproducibility, we do seed - 42 for ILP assignment + raise ValueError("seed_everything must be 42 or greater") + if not isinstance(cli.trainer.strategy, (DDPStrategy, SingleDeviceStrategy)): + # NOTE: to implement support for other distributed startegies, should check + # the places noted in this GH issue: https://github.com/StevenSong/ProtoSSL/issues/63 + raise ValueError( + f"Only single device or DDP training strategies are supported, got: {cli.trainer.strategy}" + ) + pipeline_stage: STAGE_T = cli.config.pipeline_stage + + # NOTE all model/data validation should happen in their respective modules above + # Assume at this point, we have the correct models/datasets for the given stage + if pipeline_stage == "learn-prototypes-supervised": + cli.trainer.fit( + model=cli.model, + datamodule=cli.datamodule, + ckpt_path=cli.config.resume_from_checkpoint, + ) + elif pipeline_stage == "project-prototypes-supervised": + # hijack predict over train set for prototype projection (see LitData above) + cli.datamodule.setup("fit") + cli.trainer.predict( + model=cli.model, + dataloaders=cli.datamodule.train_dataloader(), + ) + ckpt_path = os.path.join(cli.trainer.log_dir, "proj.ckpt") # type: ignore + cli.trainer.save_checkpoint(ckpt_path, weights_only=False) + elif ( + pipeline_stage == "compute-embeddings" + or pipeline_stage == "compute-fusion-embeddings" + ): + # hack to pass split name to prediction writer, not used anywhere else + cli.model.prediction_split = "train" + cli.datamodule.setup("fit") + cli.trainer.predict( + model=cli.model, + dataloaders=cli.datamodule.train_dataloader(), + ) + cli.model.prediction_split = "val" + cli.trainer.predict( + model=cli.model, + dataloaders=cli.datamodule.val_dataloader(), + ) + cli.model.prediction_split = "test" + cli.datamodule.setup("test") + cli.trainer.predict( + model=cli.model, + dataloaders=cli.datamodule.test_dataloader(), + ) + elif ( + pipeline_stage == "train-classifier" + or pipeline_stage == "train-fusion-classifier" + ): + cli.trainer.fit( + model=cli.model, + datamodule=cli.datamodule, + ckpt_path=cli.config.resume_from_checkpoint, + ) + # hack to pass split name to prediction writer, not used anywhere else + cli.model.prediction_split = "val" + cli.trainer.predict( + model=cli.model, + dataloaders=cli.datamodule.val_dataloader(), + ) + cli.model.prediction_split = "test" + cli.datamodule.setup("test") + cli.trainer.predict( + model=cli.model, + dataloaders=cli.datamodule.test_dataloader(), + ) + else: + raise ValueError(f"Unknown pipeline stage {pipeline_stage}") + check_final_link(cli.trainer.log_dir) + + +if __name__ == "__main__": + run() diff --git a/protossl/trainer.py b/protossl/trainer.py index 02c9e2a..db047ab 100644 --- a/protossl/trainer.py +++ b/protossl/trainer.py @@ -31,6 +31,7 @@ from .models import ( BaseClassifier, BlackboxClassifier, + BlackboxContraster, PrototypeAssigner, PrototypeClassifier, PrototypeContraster, @@ -54,6 +55,7 @@ def __init__( assignment_strategy: ASSIGN_T | None = None, contrastive_pair_mode: CONTRASTIVE_T | None = None, data_kwargs: dict = dict(), + data_wrapper_kwargs: dict = dict(), ): super().__init__() self.save_hyperparameters() @@ -79,6 +81,7 @@ def __init__( split="train", sampling_rate=sampling_rate, label_subset=label_subset, + **data_kwargs, ) self.label_weights = self.train_ds.get_label_weights() self.label_cooccurrence = self.train_ds.get_label_cooccurrence() @@ -94,7 +97,8 @@ def setup(self, stage: str): label_subset: list[str] | None = self.hparams.label_subset # type: ignore contrastive_pair_mode: CONTRASTIVE_T | None = self.hparams.contrastive_pair_mode # type: ignore data_kwargs: dict = self.hparams.data_kwargs # type: ignore - wrap_contrastive = pipeline_stage == "learn-prototypes" + data_wrapper_kwargs: dict = self.hparams.data_wrapper_kwargs # type: ignore + wrap_contrastive = pipeline_stage in {"learn-prototypes", "train-contraster"} if stage == "fit": if not hasattr(self, "train_ds"): @@ -116,11 +120,13 @@ def setup(self, stage: str): self.train_ds = AudioSetContrastiveWrapperDataset( self.train_ds, pair_mode=contrastive_pair_mode, - **data_kwargs, + **data_wrapper_kwargs, ) else: assert isinstance(self.train_ds, BaseTSDataset) - self.train_ds = PCLRWrapperDataset(self.train_ds, **data_kwargs) + self.train_ds = PCLRWrapperDataset( + self.train_ds, **data_wrapper_kwargs + ) if stage in ["fit", "validate"]: self.val_ds = self.ds_cls( @@ -139,11 +145,11 @@ def setup(self, stage: str): self.val_ds = AudioSetContrastiveWrapperDataset( self.val_ds, pair_mode=contrastive_pair_mode, - **data_kwargs, + **data_wrapper_kwargs, ) else: assert isinstance(self.val_ds, BaseTSDataset) - self.val_ds = PCLRWrapperDataset(self.val_ds, **data_kwargs) + self.val_ds = PCLRWrapperDataset(self.val_ds, **data_wrapper_kwargs) if stage in ["test", "predict"]: self.test_ds = self.ds_cls( @@ -318,6 +324,19 @@ def __init__( contrastive_pair_mode=contrastive_pair_mode, **model_kwargs, ) + elif pipeline_stage == "train-contraster": + if contrastive_pair_mode is None: + raise ValueError( + "pipeline_stage=train-contraster must specify contrastive_pair_mode" + ) + self.model = BlackboxContraster( + backbone_type=backbone_type, + conv_type=conv_type, + pretrained_weights=pretrained_weights, + input_channels=input_channels, + contrastive_pair_mode=contrastive_pair_mode, + **model_kwargs, + ) elif pipeline_stage == "learn-prototypes-supervised": if ( n_prototypes_per_label is None @@ -512,6 +531,30 @@ def _common_step( else: loss_terms = self.model(batch["x1"], batch["x2"]) + loss = self._log_and_composite_losses( + stage=stage, + losses=loss_terms, + batch_size=batch_size, + log=log, + sync_dist=sync_dist, + ) + elif pipeline_stage == "train-contraster": + assert isinstance(self.model, BlackboxContraster) + assert contrastive_pair_mode is not None + if stage not in ["train", "val"]: + raise ValueError( + f"Cannot use _common_step with pipeline_stage=train-contraster and (lightning) stage={stage}" + ) + + preds = None + + if contrastive_pair_mode == "cola+clar": + raise NotImplementedError( + f"BlackboxContraster does not yet support contrastive_pair_mode=cola+clar" + ) + else: + loss_terms = self.model(batch["x1"], batch["x2"]) + loss = self._log_and_composite_losses( stage=stage, losses=loss_terms, @@ -538,7 +581,6 @@ def _common_step( pipeline_stage == "project-prototypes" or pipeline_stage == "project-prototypes-supervised" ): - print(self._batch_counter, flush=True) self._batch_counter += 1 assert isinstance(self.model, PrototypeProjector) if stage != "predict": @@ -993,7 +1035,7 @@ def run(): raise ValueError("seed_everything must be 42 or greater") if not isinstance(cli.trainer.strategy, (DDPStrategy, SingleDeviceStrategy)): # NOTE: to implement support for other distributed startegies, should check - # the places noted in this GH issue: https://github.com/StevenSong/ecg-prototype-fm/issues/63 + # the places noted in this GH issue: https://github.com/StevenSong/ProtoSSL/issues/63 raise ValueError( f"Only single device or DDP training strategies are supported, got: {cli.trainer.strategy}" ) @@ -1005,6 +1047,7 @@ def run(): if ( pipeline_stage == "learn-prototypes" or pipeline_stage == "learn-prototypes-supervised" + or pipeline_stage == "train-contraster" ): cli.trainer.fit( model=cli.model, diff --git a/pyproject.toml b/pyproject.toml index 895836b..dbe0b34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ authors = [ { name="Steven Song", email="songs1@uchicago.edu" }, { name="Sahil Sethi", email="sethis@uchicago.edu" }, ] -description = "ProtoSSL: Interpretable Prototype Learning from Unlabeled Time-Series Data" +description = "ProtoSSL: Self-Supervised Pretraining and Downstream Transfer for Projection-Based Prototype Models" readme = "README.md" requires-python = ">=3.10" dynamic = ["version", "dependencies"] diff --git a/requirements.txt b/requirements.txt index ad04cfb..1f1e19a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,6 +10,7 @@ jupyterlab lightning matplotlib numpy>=2.0.0 +omegaconf optuna optuna-integration pandas diff --git a/results/audio-results.ipynb b/results/audio-results.ipynb index 01b6815..10e25bc 100644 --- a/results/audio-results.ipynb +++ b/results/audio-results.ipynb @@ -21,18 +21,18 @@ "source": [ "models = [\n", " (\"Blackbox Direct\", \"blackbox-direct\"),\n", - " (\"SupProto Direct\", \"labsup-proto-direct\"),\n", - " (\"SupProto AudioSet (Probe)\", \"labsup-proto-audioset-rila\"),\n", - " (\"SupProto AudioSet\", \"labsup-proto-audioset-rila-ft\"),\n", + " (\"SupProto Direct\", \"supproto-direct\"),\n", + " (\"SupProto AudioSet (Probe)\", \"supproto-audioset-rila\"),\n", + " (\"SupProto AudioSet\", \"supproto-audioset-rila-ft\"),\n", " (\"ProtoSSL AudioSet (Probe)\", \"protossl-audioset-pila\"),\n", " (\"ProtoSSL AudioSet\", \"protossl-audioset-pila-ft\"),\n", "]\n", "\n", "datasets = [\n", " # dataset fold path\n", - " (\"IEMOCAP\", 5, \"/opt/gpu_working/steven/protossl-audio/runs-iemocap-fold{}\"),\n", - " (\"UrbanSound8K\", 10, \"/opt/gpu_working/steven/protossl-audio/runs-us8k-fold{}\"),\n", - " (\"VoxCeleb1 ID\", None, \"/opt/gpu_working/steven/protossl-audio/runs-voxceleb\"),\n", + " (\"IEMOCAP\", 5, \"/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-iemocap-fold{}-v2\"),\n", + " (\"UrbanSound8K\", 10, \"/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-us8k-fold{}-v2\"),\n", + " (\"VoxCeleb1 ID\", None, \"/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-voxceleb-v2\"),\n", "]" ] }, diff --git a/results/ecg-results-rebuttal.ipynb b/results/ecg-results-rebuttal.ipynb new file mode 100644 index 0000000..f44603d --- /dev/null +++ b/results/ecg-results-rebuttal.ipynb @@ -0,0 +1,724 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import seaborn as sns\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Load Results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "seeds = [42, 67, 99, 70, 73]\n", + "\n", + "datasets = {\n", + " \"EchoNext\": {\n", + " 72475: \"runs-echonext\",\n", + " 32768: \"runs-echonext-32k\",\n", + " 16384: \"runs-echonext-16k\",\n", + " 8192: \"runs-echonext-8k\",\n", + " 4096: \"runs-echonext-4k\",\n", + " 2048: \"runs-echonext-2k\",\n", + " 1024: \"runs-echonext-1k\",\n", + " 512: \"runs-echonext-512\",\n", + " 256: \"runs-echonext-256\",\n", + " },\n", + " \"MIMIC-IV-ECG\": {\n", + " 78470: \"runs-mimic\",\n", + " 32768: \"runs-mimic-32k\",\n", + " 16384: \"runs-mimic-16k\",\n", + " 8192: \"runs-mimic-8k\",\n", + " 4096: \"runs-mimic-4k\",\n", + " 2048: \"runs-mimic-2k\",\n", + " 1024: \"runs-mimic-1k\",\n", + " 512: \"runs-mimic-512\",\n", + " 256: \"runs-mimic-256\",\n", + " },\n", + " \"CODE-15%\": {\n", + " 74112: \"runs-code15\",\n", + " 32768: \"runs-code15-32k\",\n", + " 16384: \"runs-code15-16k\",\n", + " 8192: \"runs-code15-8k\",\n", + " 4096: \"runs-code15-4k\",\n", + " 2048: \"runs-code15-2k\",\n", + " 1024: \"runs-code15-1k\",\n", + " 512: \"runs-code15-512\",\n", + " 256: \"runs-code15-256\",\n", + " },\n", + " \"PTB-XL\": {\n", + " 17418: \"runs-ptbxl\",\n", + " 8722: \"runs-ptbxl-8k\",\n", + " 4356: \"runs-ptbxl-4k\",\n", + " 2175: \"runs-ptbxl-2k\",\n", + " 1091: \"runs-ptbxl-1k\",\n", + " 547: \"runs-ptbxl-512\",\n", + " 273: \"runs-ptbxl-256\",\n", + " },\n", + " \"CinC Georgia\": {\n", + " 8192: \"runs-cinc\",\n", + " 4096: \"runs-cinc-4k\",\n", + " 2048: \"runs-cinc-2k\",\n", + " 1024: \"runs-cinc-1k\",\n", + " 512: \"runs-cinc-512\",\n", + " 256: \"runs-cinc-256\",\n", + " },\n", + " \"ZZU pECG\": {\n", + " 8658: \"runs-zzu\",\n", + " 4096: \"runs-zzu-4k\",\n", + " 2048: \"runs-zzu-2k\",\n", + " 1024: \"runs-zzu-1k\",\n", + " 512: \"runs-zzu-512\",\n", + " 256: \"runs-zzu-256\",\n", + " },\n", + "}\n", + "\n", + "# mapping of experiment name to tuple of:\n", + "# - plotting color\n", + "# - folder name\n", + "experiments = {\n", + " \"SupProto Direct\": (\"tab:green\", \"supproto-direct\"),\n", + " \"SupProto HEEDB\": (\"tab:orange\", \"supproto-heedb-rila\"),\n", + " \"SupProto HEEDB (FT)\": (\"tab:pink\", \"supproto-heedb-rila-ft\"),\n", + " \"ProtoSSL HEEDB\": (\"tab:blue\", \"protossl-heedb-pila\"),\n", + " \"ProtoSSL HEEDB (FT)\": (\"tab:cyan\", \"protossl-heedb-pila-ft\"),\n", + " ###\n", + " \"Blackbox Direct\": (\"tab:grey\", \"blackbox-direct\"),\n", + " \"ECGFounder\": (\"tab:red\", \"ecgfounder-logreg\"),\n", + " \"ST-MEM\": (\"tab:purple\", \"stmem-logreg\"),\n", + " \"PCLR-HEEDB\": (\"tab:olive\", \"pclr-logreg\"), # new\n", + " ###\n", + " \"SupProto HEEDB-150 (7500-proto-bank)\": (\"tab:orange\", \"supproto-heedb-150-rila\"), # new\n", + " \"SupProto HEEDB-150 (7500-proto-bank) (FT)\": (\"tab:pink\", \"supproto-heedb-150-rila-ft\"), # new\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank)\": (\"tab:blue\", \"protossl-heedb-150-pila\"), # new\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank) (FT)\": (\"tab:cyan\", \"protossl-heedb-150-pila-ft\"), # new\n", + " ###\n", + " \"SupProto HEEDB-150 (1050-proto-bank)\": (\"tab:orange\", \"supproto-heedb-150-1050-rila\"), # new\n", + " \"SupProto HEEDB-150 (1050-proto-bank) (FT)\": (\"tab:pink\", \"supproto-heedb-150-1050-rila-ft\"), # new\n", + " \"ProtoSSL HEEDB-150 (1050-proto-bank)\": (\"tab:blue\", \"protossl-heedb-150-1050-pila\"), # new\n", + " \"ProtoSSL HEEDB-150 (1050-proto-bank) (FT)\": (\"tab:cyan\", \"protossl-heedb-150-1050-pila-ft\"), # new\n", + " ###\n", + " \"ProtoSSL HEEDB (83ppl)\": (\"yellow\", \"protossl-heedb-pila-83ppl\"),\n", + " \"ProtoSSL HEEDB (PIT)\": (\"blue\", \"protossl-heedb-pit\"),\n", + " \"ProtoSSL HEEDB (PIP)\": (\"red\", \"protossl-heedb-pip\"),\n", + " \"ProtoSSL HEEDB (Latent)\": (\"orange\", \"protossl-heedb-latent\"), # new\n", + "}\n", + "\n", + "def get_palette(exp_names):\n", + " palette = dict()\n", + " exps = experiments\n", + " for exp_name in exp_names:\n", + " palette[exp_name] = exps[exp_name][0]\n", + " return palette" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "data = []\n", + "for seed in seeds:\n", + " output_dir = Path(f\"/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal/protossl-outputs-seed{seed}\")\n", + " for ds, sizes in datasets.items():\n", + " for size, run_dir in sizes.items():\n", + " exps = experiments.copy()\n", + " for exp_name, (exp_color, exp_dir) in exps.items():\n", + " _output_dir = output_dir\n", + " metrics_csv = _output_dir / run_dir / exp_dir / \"metrics-bootstrapped-v2.csv\"\n", + " if not os.path.exists(metrics_csv):\n", + " continue\n", + " metrics = pd.read_csv(metrics_csv, index_col=\"Label\")\n", + " multilabel = metrics.loc[\"Multilabel Averaged\"]\n", + " datum = {\n", + " \"Seed\": seed,\n", + " \"Dataset\": ds,\n", + " \"Model\": exp_name,\n", + " \"Train Size\": size,\n", + " \"Multilabel (AUROC)\": multilabel[\"AUROC\"],\n", + " \"Multilabel (AUPRC)\": multilabel[\"AUPRC\"],\n", + " }\n", + " if \"AUROC 95% CI (lo)\" in metrics.columns:\n", + " datum[\"AUROC 95% CI (lo)\"] = metrics.loc[\"Multilabel Averaged\", \"AUROC 95% CI (lo)\"]\n", + " datum[\"AUROC 95% CI (hi)\"] = metrics.loc[\"Multilabel Averaged\", \"AUROC 95% CI (hi)\"]\n", + " data.append(datum)\n", + "results = pd.DataFrame.from_records(data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "len(results)" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## Plot Label Efficiency Curve" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_lift(\n", + " *, # enforce kwargs\n", + " df: pd.DataFrame, # long format df (each point to plot is a row)\n", + " dataset: str,\n", + " seed: int | list[int] = 42, # if list of int, average across seeds\n", + " metric: str,\n", + " models: list[str], # must be intentional about which models to plot\n", + " rename: list[str] | None = None,\n", + " baseline_model: str | None = None, # singular result to optionally plot as dashed line\n", + " ylim: tuple[float, float] | None = None,\n", + " xlim: tuple[float, float] | None = None,\n", + " save_path: str | None = None,\n", + " smooth: bool = True,\n", + " ax: plt.Axes | None = None,\n", + " do_label: bool = True,\n", + ") -> tuple[plt.Figure, plt.Axes]:\n", + " if rename is not None:\n", + " assert len(models) == len(rename)\n", + " if baseline_model is not None and baseline_model not in models:\n", + " models = [baseline_model] + models\n", + " if rename is not None:\n", + " rename = [baseline_model] + rename\n", + " palette = get_palette(models)\n", + " if rename is not None:\n", + " palette = {new_name: palette[m] for new_name, m in zip(rename, models)}\n", + " df = df.copy()\n", + " df[\"Model\"] = df[\"Model\"].replace({v: k for k, v in zip(rename, models)})\n", + "\n", + " df = df[df[\"Dataset\"] == dataset]\n", + " if isinstance(seed, int):\n", + " df = df[df[\"Seed\"] == seed]\n", + " else: # list of seeds\n", + " df = df[df[\"Seed\"].isin(seed)]\n", + " df = df.groupby([\"Dataset\", \"Model\", \"Train Size\"]).mean().reset_index()\n", + " # print(dataset, df[metric].min(), df[metric].max())\n", + " fig = None\n", + " if ax is None:\n", + " fig, ax = plt.subplots(figsize=(6, 6))\n", + " min_size = df[\"Train Size\"].min()\n", + " max_size = df[\"Train Size\"].max()\n", + " if xlim is not None:\n", + " min_size = min(min_size, xlim[0])\n", + " max_size = max(max_size, xlim[1])\n", + "\n", + " if baseline_model is not None:\n", + " mask = df[\"Model\"] == baseline_model\n", + " assert (\n", + " mask.sum() == 1\n", + " ), f\"Should only have 1 entry for baseline model: {baseline_model}\"\n", + " baseline_row = df[mask].iloc[0]\n", + " ax.hlines(\n", + " baseline_row[metric],\n", + " min_size,\n", + " max_size,\n", + " colors=palette.pop(baseline_model),\n", + " linestyles=\":\",\n", + " label=baseline_model,\n", + " )\n", + " df = df[~mask] # subsequent line plots should exclude baseline model\n", + "\n", + " if not smooth:\n", + " sns.lineplot(\n", + " df,\n", + " x=\"Train Size\",\n", + " y=metric,\n", + " hue=\"Model\",\n", + " palette=palette,\n", + " hue_order=list(palette.keys()),\n", + " marker=\"o\",\n", + " ax=ax,\n", + " )\n", + " else:\n", + " for model, color in palette.items():\n", + " model_data = df[df[\"Model\"] == model]\n", + " ax.set_xlim([2**7, 2**17])\n", + " sns.regplot(\n", + " model_data,\n", + " x=\"Train Size\",\n", + " y=metric,\n", + " color=color,\n", + " marker=\"o\",\n", + " ax=ax,\n", + " logx=True,\n", + " label=model if do_label else None,\n", + " line_kws={\"zorder\": 2},\n", + " scatter_kws={\"zorder\": 3},\n", + " truncate=False,\n", + " )\n", + " ax.set_xscale(\"log\", base=2)\n", + " if xlim is not None:\n", + " ax.set_xlim(xlim)\n", + " else:\n", + " # tighter boundaries than default lims\n", + " ax.set_xlim((min_size, max_size))\n", + " ax.set_title(dataset)\n", + " # ax.set_title(f\"{dataset} {metric}\")\n", + " if do_label:\n", + " ax.legend(loc=\"lower right\")\n", + " if ylim is not None:\n", + " ax.set_ylim(ylim)\n", + " else:\n", + " ymin, ymax = ax.get_ylim()\n", + " ax.set_ylim((ymin, min(ymax, 1)))\n", + " if save_path is not None:\n", + " assert fig is not None\n", + " fig.tight_layout()\n", + " fig.savefig(save_path)\n", + " return fig, ax" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "### Figs to help visualize tables" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plot_lift(\n", + " df=results,\n", + " dataset=\"EchoNext\",\n", + " metric=\"Multilabel (AUROC)\",\n", + " models=[\n", + " \"SupProto Direct\",\n", + " \"SupProto HEEDB-150 (7500-proto-bank)\",\n", + " \"SupProto HEEDB-150 (7500-proto-bank) (FT)\",\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank)\",\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank) (FT)\",\n", + " ],\n", + " seed=seeds,\n", + " do_label=True,\n", + ")\n", + "ax.set_title(\"EchoNext (7500 Pretrained Prototypes, 14 PPL Downstream)\")\n", + "ax.set_ylim([0.5, 0.9])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plot_lift(\n", + " df=results,\n", + " dataset=\"EchoNext\",\n", + " metric=\"Multilabel (AUROC)\",\n", + " models=[\n", + " \"SupProto Direct\",\n", + " \"SupProto HEEDB-150 (1050-proto-bank)\",\n", + " \"SupProto HEEDB-150 (1050-proto-bank) (FT)\",\n", + " \"ProtoSSL HEEDB-150 (1050-proto-bank)\",\n", + " \"ProtoSSL HEEDB-150 (1050-proto-bank) (FT)\",\n", + " ],\n", + " seed=seeds,\n", + " do_label=True,\n", + ")\n", + "ax.set_title(\"EchoNext (1050 Pretrained Prototypes, 14 PPL Downstream)\")\n", + "ax.set_ylim([0.5, 0.9])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plot_lift(\n", + " df=results,\n", + " dataset=\"EchoNext\",\n", + " metric=\"Multilabel (AUROC)\",\n", + " models=[\n", + " \"Blackbox Direct\",\n", + " \"ProtoSSL HEEDB\",\n", + " \"ECGFounder\",\n", + " \"ST-MEM\",\n", + " \"PCLR-HEEDB\",\n", + " ],\n", + " seed=seeds,\n", + " do_label=True,\n", + ")\n", + "ax.set_title(\"EchoNext\")\n", + "ax.set_ylim([0.5, 0.9])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plot_lift(\n", + " df=results,\n", + " dataset=\"EchoNext\",\n", + " metric=\"Multilabel (AUROC)\",\n", + " models=[\n", + " \"ProtoSSL HEEDB (83ppl)\",\n", + " \"ProtoSSL HEEDB (PIT)\",\n", + " \"ProtoSSL HEEDB (PIP)\",\n", + " \"ProtoSSL HEEDB (Latent)\",\n", + " ],\n", + " seed=seeds,\n", + " do_label=True,\n", + ")\n", + "ax.set_title(\"EchoNext\")\n", + "ax.set_ylim([0.5, 0.9])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "print(results[(results[\"Model\"] == \"PCLR-HEEDB\") & (results[\"Train Size\"] == 72475)].to_markdown())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "print(results[(results[\"Model\"] == \"ProtoSSL HEEDB (83ppl)\") & (results[\"Train Size\"] == 72475)].to_markdown())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plot_lift(\n", + " df=results,\n", + " dataset=\"EchoNext\",\n", + " metric=\"Multilabel (AUROC)\",\n", + " models=[\n", + " \"ProtoSSL HEEDB (83ppl)\",\n", + " \"PCLR-HEEDB\",\n", + " ],\n", + " seed=seeds,\n", + " do_label=True,\n", + ")\n", + "ax.set_title(\"EchoNext\")\n", + "ax.set_ylim([0.5, 0.9])" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "## Make Tables" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "all_agg = (\n", + " results\n", + " .sort_values([\"Model\", \"Dataset\", \"Train Size\"])\n", + " .groupby([\"Dataset\", \"Model\", \"Train Size\"])\n", + " # [\"Multilabel (AUROC)\"].describe() # reporting std dev\n", + " [[\"Multilabel (AUROC)\", \"AUROC 95% CI (lo)\", \"AUROC 95% CI (hi)\"]].mean() # reporting bootstrapped CI\n", + ")\n", + "# all_agg = all_agg.reset_index()[[\"Dataset\", \"Model\", \"Train Size\", \"mean\", \"std\"]] # reporting std dev\n", + "all_agg = all_agg.reset_index()[[\"Dataset\", \"Model\", \"Train Size\", \"Multilabel (AUROC)\", \"AUROC 95% CI (lo)\", \"AUROC 95% CI (hi)\"]] # reporting bootstrapped CI\n", + "\n", + "# filt_agg = all_agg[\n", + "# all_agg[\"Model\"].isin([\n", + "# \"SupProto Direct\",\n", + "# \"SupProto HEEDB\",\n", + "# \"SupProto HEEDB (FT)\",\n", + "# \"ProtoSSL HEEDB\",\n", + "# \"ProtoSSL HEEDB (FT)\",\n", + "# ])\n", + "# ].reset_index(drop=True)\n", + "filt_agg = all_agg\n", + "\n", + "# reporting std dev\n", + "# filt_agg[\"val_mean\"] = filt_agg[\"mean\"].apply(lambda x: f\"{x:0.3f}\")\n", + "# filt_agg[\"val_std\"] = filt_agg[\"std\"].apply(lambda x: f\"{x:.2e}\")\n", + "# filt_agg[\"val\"] = filt_agg[\"val_mean\"] + \" ± \" + filt_agg[\"val_std\"]\n", + "\n", + "# report bootstrapped CI\n", + "filt_agg[\"val\"] = (\n", + " filt_agg[\"Multilabel (AUROC)\"].apply(lambda x: f\"{x:0.3f}\")\n", + " + \" [\"\n", + " + filt_agg[\"AUROC 95% CI (lo)\"].apply(lambda x: f\"{x:0.3f}\")\n", + " + \"-\"\n", + " + filt_agg[\"AUROC 95% CI (hi)\"].apply(lambda x: f\"{x:0.3f}\")\n", + " + \"]\"\n", + ")\n", + "\n", + "pivoted = filt_agg[[\"Dataset\", \"Model\", \"Train Size\", \"val\"]].pivot(columns=[\"Dataset\", \"Train Size\"], index=[\"Model\"], values=\"val\")\n", + "pivoted = pivoted.sort_index(axis=1, level=[0, 1], ascending=[True, False])\n", + "pivoted.index.name = None" + ] + }, + { + "cell_type": "markdown", + "id": "17", + "metadata": {}, + "source": [ + "### HEEDB 150 Labels (new)" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "#### 1050 Pretrained Prototypes (7 PPL during pretraining for SupProto HEEDB)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "n1050 = pivoted.loc[\n", + " [\"ProtoSSL HEEDB-150 (1050-proto-bank) (FT)\", \"SupProto HEEDB-150 (1050-proto-bank) (FT)\", \"ProtoSSL HEEDB-150 (1050-proto-bank)\", \"SupProto HEEDB-150 (1050-proto-bank)\", \"SupProto Direct\"],\n", + " [\"EchoNext\"]\n", + " ].T.copy().rename(columns=dict(zip(\n", + " [\"ProtoSSL HEEDB-150 (1050-proto-bank) (FT)\", \"SupProto HEEDB-150 (1050-proto-bank) (FT)\", \"ProtoSSL HEEDB-150 (1050-proto-bank)\", \"SupProto HEEDB-150 (1050-proto-bank)\", \"SupProto Direct\"],\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"]\n", + " )))\n", + "\n", + "print(n1050.to_markdown())" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "#### 7500 Pretrained Prototypes (50 PPL during pretraining for SupProto HEEDB)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "n7500 = pivoted.loc[\n", + " [\"ProtoSSL HEEDB-150 (7500-proto-bank) (FT)\", \"SupProto HEEDB-150 (7500-proto-bank) (FT)\", \"ProtoSSL HEEDB-150 (7500-proto-bank)\", \"SupProto HEEDB-150 (7500-proto-bank)\", \"SupProto Direct\"],\n", + " [\"EchoNext\"]\n", + " ].T.copy().rename(columns=dict(zip(\n", + " [\"ProtoSSL HEEDB-150 (7500-proto-bank) (FT)\", \"SupProto HEEDB-150 (7500-proto-bank) (FT)\", \"ProtoSSL HEEDB-150 (7500-proto-bank)\", \"SupProto HEEDB-150 (7500-proto-bank)\", \"SupProto Direct\"],\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"]\n", + " )))\n", + "\n", + "print(n7500.to_markdown())" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "### Blackbox Baselines - including PCLR (new)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "pivoted.loc[\n", + " [\"PCLR-HEEDB\", \"ProtoSSL HEEDB\", \"ECGFounder\", \"ST-MEM\", \"Blackbox Direct\"],\n", + " [\"EchoNext\", \"MIMIC-IV-ECG\", \"ZZU pECG\", \"PTB-XL\", \"CinC Georgia\", \"CODE-15%\"]\n", + "].T" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"PCLR-HEEDB\", \"ProtoSSL HEEDB\", \"ECGFounder\", \"ST-MEM\", \"Blackbox Direct\"],\n", + " [\"EchoNext\"]\n", + " ].T.to_markdown()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "25", + "metadata": {}, + "source": [ + "### No Assignment (PIT & PIP) - and No Projection (Latent) (new)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB (83ppl)\", \"ProtoSSL HEEDB (PIT)\", \"ProtoSSL HEEDB (PIP)\", \"ProtoSSL HEEDB (Latent)\"],\n", + " [\"EchoNext\"]\n", + " ].T.rename(columns={\n", + " \"ProtoSSL HEEDB (83ppl)\": \"ProtoSSL HEEDB (LAP) (83PPL)\",\n", + " }).to_markdown()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "27", + "metadata": {}, + "source": [ + "### ZZU full scale more prototypes in ProtoSSL pretraining" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB\", \"ProtoSSL HEEDB-150 (7500-proto-bank)\", \"SupProto Direct\"],\n", + " [\"ZZU pECG\"]\n", + " ].T.to_markdown()\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB\", \"ProtoSSL HEEDB-150 (7500-proto-bank)\", \"SupProto Direct\"],\n", + " [\"ZZU pECG\"]\n", + " ].T.rename(columns={\n", + " \"ProtoSSL HEEDB\": \"ProtoSSL HEEDB (1000-proto-bank)\",\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank)\": \"ProtoSSL HEEDB (7500 proto-bank)\"\n", + " }).to_markdown()\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "temp = filt_agg[[\"Dataset\", \"Model\", \"Train Size\", \"Multilabel (AUROC)\"]].pivot(columns=[\"Dataset\", \"Train Size\"], index=[\"Model\"], values=\"Multilabel (AUROC)\")\n", + "temp = temp.sort_index(axis=1, level=[0, 1], ascending=[True, False])\n", + "temp.index.name = None\n", + "print(\n", + " temp.loc[\n", + " [\"ProtoSSL HEEDB\", \"ProtoSSL HEEDB-150 (7500-proto-bank)\", \"SupProto Direct\"],\n", + " [\"ZZU pECG\"]\n", + " ].T.rename(columns={\n", + " \"ProtoSSL HEEDB\": \"ProtoSSL HEEDB (1000-proto-bank)\",\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank)\": \"ProtoSSL HEEDB (7500 proto-bank)\"\n", + " }).to_markdown()\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "protossl", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/results/ecg-results.ipynb b/results/ecg-results.ipynb index 2ec018b..1307a48 100644 --- a/results/ecg-results.ipynb +++ b/results/ecg-results.ipynb @@ -98,9 +98,9 @@ "# - plotting color\n", "# - folder name\n", "experiments = {\n", - " \"SupProto Direct\": (\"tab:green\", \"labsup-proto-direct\"),\n", - " \"SupProto HEEDB\": (\"tab:orange\", \"labsup-proto-heedb-rila\"),\n", - " \"SupProto HEEDB (FT)\": (\"tab:pink\", \"labsup-proto-heedb-rila-ft\"),\n", + " \"SupProto Direct\": (\"tab:green\", \"supproto-direct\"),\n", + " \"SupProto HEEDB\": (\"tab:orange\", \"supproto-heedb-rila\"),\n", + " \"SupProto HEEDB (FT)\": (\"tab:pink\", \"supproto-heedb-rila-ft\"),\n", " \"ProtoSSL HEEDB\": (\"tab:blue\", \"protossl-heedb-pila\"),\n", " \"ProtoSSL HEEDB (FT)\": (\"tab:cyan\", \"protossl-heedb-pila-ft\"),\n", " ###\n", @@ -108,18 +108,18 @@ " \"ECGFounder\": (\"tab:red\", \"ecgfounder-logreg\"),\n", " \"ST-MEM\": (\"tab:purple\", \"stmem-logreg\"),\n", " ###\n", - " \"SupProto HEEDB (PILA)\": (\"tab:red\", \"labsup-proto-heedb-pila\"),\n", - " \"SupProto HEEDB (PILA) (FT)\": (\"tab:grey\", \"labsup-proto-heedb-pila-ft\"),\n", + " \"SupProto HEEDB (PILA)\": (\"tab:red\", \"supproto-heedb-pila\"),\n", + " \"SupProto HEEDB (PILA) (FT)\": (\"tab:grey\", \"supproto-heedb-pila-ft\"),\n", " ###\n", - " \"SupProto Direct (7ppl)\": (\"tab:green\", \"labsup-proto-direct-7ppl\"),\n", - " \"SupProto HEEDB (7ppl)\": (\"tab:orange\", \"labsup-proto-heedb-rila-7ppl\"),\n", - " \"SupProto HEEDB (7ppl) (FT)\": (\"tab:pink\", \"labsup-proto-heedb-rila-7ppl-ft\"),\n", + " \"SupProto Direct (7ppl)\": (\"tab:green\", \"supproto-direct-7ppl\"),\n", + " \"SupProto HEEDB (7ppl)\": (\"tab:orange\", \"supproto-heedb-rila-7ppl\"),\n", + " \"SupProto HEEDB (7ppl) (FT)\": (\"tab:pink\", \"supproto-heedb-rila-7ppl-ft\"),\n", " \"ProtoSSL HEEDB (7ppl)\": (\"tab:blue\", \"protossl-heedb-pila-7ppl\"),\n", " \"ProtoSSL HEEDB (7ppl) (FT)\": (\"tab:cyan\", \"protossl-heedb-pila-7ppl-ft\"),\n", " ###\n", - " \"SupProto Direct (28ppl)\": (\"tab:green\", \"labsup-proto-direct-28ppl\"),\n", - " \"SupProto HEEDB (28ppl)\": (\"tab:orange\", \"labsup-proto-heedb-rila-28ppl\"),\n", - " \"SupProto HEEDB (28ppl) (FT)\": (\"tab:pink\", \"labsup-proto-heedb-rila-28ppl-ft\"),\n", + " \"SupProto Direct (28ppl)\": (\"tab:green\", \"supproto-direct-28ppl\"),\n", + " \"SupProto HEEDB (28ppl)\": (\"tab:orange\", \"supproto-heedb-rila-28ppl\"),\n", + " \"SupProto HEEDB (28ppl) (FT)\": (\"tab:pink\", \"supproto-heedb-rila-28ppl-ft\"),\n", " \"ProtoSSL HEEDB (28ppl)\": (\"tab:blue\", \"protossl-heedb-pila-28ppl\"),\n", " \"ProtoSSL HEEDB (28ppl) (FT)\": (\"tab:cyan\", \"protossl-heedb-pila-28ppl-ft\"),\n", " ###\n", @@ -136,33 +136,6 @@ " \"ECGFounder (Rand)\": (\"maroon\", \"ecgfounder-random-1000-init\"),\n", "}\n", "\n", - "use_new = [\n", - " \"Blackbox Direct\",\n", - " \"SupProto Direct\",\n", - " \"SupProto HEEDB\",\n", - " \"SupProto HEEDB (FT)\",\n", - " \"SupProto HEEDB (PILA)\",\n", - " \"SupProto HEEDB (PILA) (FT)\",\n", - " \"SupProto Direct (7ppl)\",\n", - " \"SupProto HEEDB (7ppl)\",\n", - " \"SupProto HEEDB (7ppl) (FT)\",\n", - " \"SupProto Direct (28ppl)\",\n", - " \"SupProto HEEDB (28ppl)\",\n", - " \"SupProto HEEDB (28ppl) (FT)\",\n", - " \"ProtoSSL HEEDB\",\n", - " \"ProtoSSL HEEDB (FT)\",\n", - " \"ProtoSSL HEEDB (PIA)\",\n", - " \"ProtoSSL HEEDB (PIA) (FT)\", # new\n", - " \"ProtoSSL HEEDB (PIT)\",\n", - " \"ProtoSSL HEEDB (PIP)\",\n", - " \"ProtoSSL HEEDB (7ppl)\",\n", - " \"ProtoSSL HEEDB (7ppl) (FT)\",\n", - " \"ProtoSSL HEEDB (28ppl)\",\n", - " \"ProtoSSL HEEDB (28ppl) (FT)\",\n", - " \"ProtoSSL HEEDB (83ppl)\",\n", - " \"ProtoSSL HEEDB (PIA) (83ppl)\",\n", - "]\n", - "\n", "def get_palette(exp_names):\n", " palette = dict()\n", " exps = experiments\n", @@ -180,16 +153,12 @@ "source": [ "data = []\n", "for seed in seeds:\n", - " output_dir = Path(f\"/opt/gpu_working/steven/protossl-outputs-seed{seed}\")\n", + " output_dir = Path(f\"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed{seed}\")\n", " for ds, sizes in datasets.items():\n", " for size, run_dir in sizes.items():\n", " exps = experiments.copy()\n", " for exp_name, (exp_color, exp_dir) in exps.items():\n", - " if exp_name in use_new:\n", - " _output_dir = Path(f\"/opt/gpu_working/steven/new/protossl-outputs-seed{seed}\")\n", - " else:\n", - " _output_dir = output_dir\n", - " metrics_csv = _output_dir / run_dir / exp_dir / \"metrics-bootstrapped-v2.csv\"\n", + " metrics_csv = output_dir / run_dir / exp_dir / \"metrics-bootstrapped-v2.csv\"\n", " if not os.path.exists(metrics_csv):\n", " continue\n", " metrics = pd.read_csv(metrics_csv, index_col=\"Label\")\n", diff --git a/results/meta-results.ipynb b/results/meta-results.ipynb index 0ca3815..eb9c9f5 100644 --- a/results/meta-results.ipynb +++ b/results/meta-results.ipynb @@ -53,8 +53,8 @@ "source": [ "experiments = {\n", " # nodes we ran on have different relative speeds so only use runs that entirely executed on a single node\n", - " \"ProtoSSL HEEDB\": \"/opt/gpu_working/steven/new/protossl-outputs-seed{seed}/{run_dir}/protossl-heedb-pila\",\n", - " \"ProtoSSL HEEDB (PIA)\": \"/opt/gpu_working/steven/new/protossl-outputs-seed{seed}/{run_dir}/protossl-heedb-pia\",\n", + " \"ProtoSSL HEEDB\": \"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed{seed}/{run_dir}/protossl-heedb-pila\",\n", + " \"ProtoSSL HEEDB (PIA)\": \"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed{seed}/{run_dir}/protossl-heedb-pia\",\n", "}\n", "\n", "data = []\n", @@ -120,14 +120,14 @@ "source": [ "experiments = {\n", " # nodes we ran on have different relative speeds so only use runs that entirely executed on a single node\n", - " \"ProtoSSL HEEDB\": \"/opt/gpu_working/steven/protossl-outputs-seed{seed}/{run_dir}/protossl-heedb-pila\",\n", - " \"ProtoSSL HEEDB (PIA)\": \"/opt/gpu_working/steven/protossl-outputs-seed{seed}/{run_dir}/protossl-heedb-pia\",\n", + " \"ProtoSSL HEEDB\": \"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed{seed}/{run_dir}/protossl-heedb-pila\",\n", + " \"ProtoSSL HEEDB (PIA)\": \"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed{seed}/{run_dir}/protossl-heedb-pia\",\n", "}\n", "\n", "data = []\n", "for subset_size, run_dir in echonext_dirs.items():\n", " for seed in [42, 67, 70, 73, 99]:\n", - " model_dir = Path(f\"/opt/gpu_working/steven/protossl-outputs-seed{seed}/{run_dir}/ecgfounder-clustering\")\n", + " model_dir = Path(f\"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed{seed}/{run_dir}/ecgfounder-clustering\")\n", " temp = np.load(model_dir / \"assign_frec.npy\")\n", "\n", " datum = {\n", @@ -170,7 +170,7 @@ "for seed in [42, 67, 70, 73, 99]:\n", " for subset_size, run_dir in echonext_dirs.items():\n", " models = joblib.load(\n", - " f\"/opt/gpu_working/steven/new/protossl-outputs-seed{seed}/{run_dir}/protossl-heedb-pila/model.joblib\"\n", + " f\"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed{seed}/{run_dir}/protossl-heedb-pila/model.joblib\"\n", " )\n", "\n", " K = 14\n", diff --git a/results/proto-recur-multilabel.ipynb b/results/proto-recur-multilabel.ipynb new file mode 100644 index 0000000..9ea4c68 --- /dev/null +++ b/results/proto-recur-multilabel.ipynb @@ -0,0 +1,865 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Prototype recurrence across downstream datasets\n", + "\n", + "Prototypes are learned once over HEEDB during pretraining. Each downstream dataset\n", + "(PTB-XL, CinC, CODE-15%) then independently *assigns* a fixed number of those\n", + "prototypes to each of its labels via the ILP step. A prototype that several datasets\n", + "pick for the same clinical concept is **recurrent**: the same latent prototype is doing\n", + "the same work everywhere, which is the evidence that the prototype captures the\n", + "concept rather than a dataset artifact.\n", + "\n", + "This notebook operates over multiple labels and produces:\n", + "\n", + "1. per-label overlap tables and Venn diagrams of the assigned prototype sets,\n", + "2. the list of recurrent prototypes per label,\n", + "3. clinical printouts of a chosen prototype projected onto real examples under:\n", + " * label-free projection in HEEDB\n", + " * label-supervised projection in downstream settings (both without and with fine-tuning)\n", + "\n", + "Everything below is driven by `LABELS` (the label mapping) and `DATASETS` (the run\n", + "registry), so adding a label or a dataset only means editing those two constants." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"CACHE_DIR\"] = \"/home/songs1/.cache/protossl_cache\" # for working in dev container" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import itertools\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "\n", + "from protossl.datasets import (\n", + " HeedbECGDataset,\n", + " CincECGDataset,\n", + " Code15ECGDataset,\n", + " PtbxlECGDataset,\n", + ")\n", + "from protossl.defines import (\n", + " HEEDB_TARGETS,\n", + " CINC_TARGETS,\n", + " CODE15_TARGETS,\n", + " PTBXL_TARGETS,\n", + ")\n", + "from protossl.plotting import plot_ecg\n", + "\n", + "SEED = 42\n", + "\n", + "RUN_DIR = Path(\"/opt/gpu_working/steven/protossl-ecg-outputs\")\n", + "PRETRAIN_DIR = RUN_DIR / \"protossl-heedb\"\n", + "EXPERIMENT_DIR = RUN_DIR / f\"experiments-seed{SEED}\"\n", + "MODEL = \"protossl-heedb-pila\"\n", + "PPL = 14\n", + "\n", + "DATA_DIR = Path(\"/opt/gpudata/ecg\")\n", + "SAMPLING_RATE = 100\n", + "SPLIT = \"train\"\n", + "\n", + "FIG_DIR = Path(\"figs\")\n", + "FIG_DIR.mkdir(parents=True, exist_ok=True)" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## Datasets and labels\n", + "\n", + "`DATASETS` holds the three downstream datasets that assign prototypes; HEEDB is kept\n", + "separate in `PRETRAIN` because it is the source of the prototypes, not a consumer of\n", + "them (its pretraining is unsupervised, so it never assigns prototypes to labels).\n", + "\n", + "`LABELS` maps a canonical label name onto the exact target string used by each\n", + "dataset's label set. Edit the right-hand sides to change the clinical mapping; the\n", + "next cell checks every string against the corresponding `*_TARGETS` in\n", + "`protossl.defines`. Two mappings are judgement calls worth knowing about:\n", + "\n", + "- **RBBB**: PTB-XL and CinC distinguish complete from incomplete RBBB, so the mapping\n", + " uses `CRBBB`. CinC also carries a plain `RBBB` code; swap it in if the looser\n", + " definition is wanted (`IRBBB` is the incomplete-only code in both).\n", + "- **LBBB**: PTB-XL again splits complete/incomplete, so the mapping uses `CLBBB`.\n", + " CinC has only `LBBB`/`ILBBB`, so `LBBB` there already means complete." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "PRETRAIN = {\n", + " \"slug\": \"heedb\",\n", + " \"dataset_cls\": HeedbECGDataset,\n", + " \"dataset_path\": DATA_DIR / \"heedb\",\n", + " \"targets\": list(HEEDB_TARGETS),\n", + "}\n", + "\n", + "DATASETS = {\n", + " \"PTB-XL\": {\n", + " \"slug\": \"ptbxl\",\n", + " \"run_dir\": EXPERIMENT_DIR / \"runs-ptbxl\",\n", + " \"dataset_cls\": PtbxlECGDataset,\n", + " \"dataset_path\": DATA_DIR / \"ptb-xl\",\n", + " \"targets\": PTBXL_TARGETS,\n", + " },\n", + " \"CinC\": {\n", + " \"slug\": \"cinc\",\n", + " \"run_dir\": EXPERIMENT_DIR / \"runs-cinc\",\n", + " \"dataset_cls\": CincECGDataset,\n", + " \"dataset_path\": DATA_DIR / \"cinc-2020\",\n", + " \"targets\": CINC_TARGETS,\n", + " },\n", + " \"CODE-15%\": {\n", + " \"slug\": \"code15\",\n", + " \"run_dir\": EXPERIMENT_DIR / \"runs-code15\",\n", + " \"dataset_cls\": Code15ECGDataset,\n", + " \"dataset_path\": DATA_DIR / \"code15\",\n", + " \"targets\": CODE15_TARGETS,\n", + " },\n", + "}\n", + "\n", + "# canonical label -> that label's exact target string in each label set\n", + "LABELS = {\n", + " \"1st-degree AV block\": {\n", + " \"HEEDB\": \"WITH 1ST DEGREE AV BLOCK\",\n", + " \"PTB-XL\": \"1AVB\",\n", + " \"CinC\": \"IAVB\",\n", + " \"CODE-15%\": \"1dAVb\",\n", + " },\n", + " \"RBBB\": {\n", + " \"HEEDB\": \"RIGHT BUNDLE BRANCH BLOCK\",\n", + " \"PTB-XL\": \"CRBBB\",\n", + " \"CinC\": \"CRBBB\",\n", + " \"CODE-15%\": \"RBBB\",\n", + " },\n", + " \"LBBB\": {\n", + " \"HEEDB\": \"LEFT BUNDLE BRANCH BLOCK\",\n", + " \"PTB-XL\": \"CLBBB\",\n", + " \"CinC\": \"LBBB\",\n", + " \"CODE-15%\": \"LBBB\",\n", + " },\n", + " \"AFib\": {\n", + " \"HEEDB\": \"ATRIAL FIBRILLATION\",\n", + " \"PTB-XL\": \"AFIB\",\n", + " \"CinC\": \"AF\",\n", + " \"CODE-15%\": \"AF\",\n", + " },\n", + " \"Sinus Brady\": {\n", + " \"HEEDB\": \"SINUS BRADYCARDIA\",\n", + " \"PTB-XL\": \"SBRAD\",\n", + " \"CinC\": \"SB\",\n", + " \"CODE-15%\": \"SB\",\n", + " },\n", + " \"Sinus Tach\": {\n", + " \"HEEDB\": \"SINUS TACHYCARDIA\",\n", + " \"PTB-XL\": \"STACH\",\n", + " \"CinC\": \"STach\",\n", + " \"CODE-15%\": \"ST\",\n", + " },\n", + "}\n", + "\n", + "# filename-safe stems for saved figures\n", + "LABEL_SLUGS = {\n", + " \"1st-degree AV block\": \"1davb\",\n", + " \"RBBB\": \"rbbb\",\n", + " \"LBBB\": \"lbbb\",\n", + " \"AFib\": \"afib\",\n", + " \"Sinus Brady\": \"sbrad\",\n", + " \"Sinus Tach\": \"stach\",\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "# every mapped name must exist in the label set it refers to, otherwise a typo\n", + "# silently turns into an empty prototype set downstream\n", + "for label, mapping in LABELS.items():\n", + " assert label in LABEL_SLUGS, f\"no filename slug for {label}\"\n", + " assert mapping[\"HEEDB\"] in PRETRAIN[\"targets\"], f\"{label}: bad HEEDB target\"\n", + " for name, cfg in DATASETS.items():\n", + " assert mapping[name] in cfg[\"targets\"], f\"{label}: bad {name} target\"\n", + "\n", + "pd.DataFrame(LABELS).T" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## Load assignment and projection metadata\n", + "\n", + "- `assignment_metadata.csv` (one row per label slot) says which pretrained prototype\n", + " the ILP gave to each label slot of a downstream dataset.\n", + "- `projection_metadata.csv` (one row per prototype) says which sample and 1-second\n", + " window each prototype was projected onto, i.e. the waveform to plot." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "def read_metadata(path: Path) -> pd.DataFrame:\n", + " \"\"\"Read a metadata CSV, normalizing the sample ID column name.\n", + "\n", + " `ecg_id` was renamed to `sample_id` when the datasets were generalized beyond\n", + " ECG (see `protossl.trainer.PredictionWriter`), so runs written before that\n", + " rename still carry the old column name. Same with `patient_id` --> `source_id.\n", + " \"\"\"\n", + " return pd.read_csv(path).rename(columns={\"ecg_id\": \"sample_id\", \"patient_id\": \"source_id\"})\n", + "\n", + "\n", + "heedb_proj = read_metadata(\n", + " PRETRAIN_DIR / \"project-prototypes/latest/projection_metadata.csv\"\n", + ")\n", + "\n", + "for name, cfg in DATASETS.items():\n", + " run = cfg[\"run_dir\"]\n", + " cfg[\"assign\"] = read_metadata(\n", + " run / MODEL / \"learn-prototype-assignments/latest/assignment_metadata.csv\"\n", + " )\n", + " cfg[\"proj\"] = read_metadata(\n", + " run / MODEL / \"project-prototypes-supervised/latest/projection_metadata.csv\"\n", + " )\n", + " cfg[\"ft_proj\"] = read_metadata(\n", + " run / f\"{MODEL}-ft\" / \"project-prototypes-supervised/latest/projection_metadata.csv\"\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "# the row position of a prototype in `proj`/`ft_proj` is its slot in the projected\n", + "# model, and assignment rows are written label-major then slot-minor, so an\n", + "# assignment row's position is exactly that slot. the checks below pin down that\n", + "# correspondence, which is what lets us index projections by assignment row.\n", + "assert (heedb_proj[\"prototype_id\"] == heedb_proj.index).all()\n", + "\n", + "for name, cfg in DATASETS.items():\n", + " assign, proj, ft_proj = cfg[\"assign\"], cfg[\"proj\"], cfg[\"ft_proj\"]\n", + " n_slots = len(cfg[\"targets\"]) * PPL\n", + " assert len(assign) == len(proj) == len(ft_proj) == n_slots, name\n", + " assert (assign.index == assign[\"label_idx\"] * PPL + assign[\"slot_idx\"]).all(), name\n", + " assert (proj[\"prototype_id\"] == proj.index).all(), name\n", + " assert (ft_proj[\"prototype_id\"] == ft_proj.index).all(), name\n", + " # prototypes come from the pretrained bank, which HEEDB projection enumerates\n", + " assert assign[\"prototype_idx\"].isin(heedb_proj.index).all(), name\n", + "\n", + "print(f\"{len(heedb_proj)} pretrained prototypes\")" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "## Prototypes assigned to each label\n", + "\n", + "The ILP (`ilp_effect_size`) uses each prototype at most once per dataset, so within a\n", + "dataset the prototype sets of different labels are disjoint, and each label gets\n", + "exactly `PPL` prototypes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "def label_prototypes(label: str) -> dict[str, set[int]]:\n", + " \"\"\"Pretrained prototype indices each dataset assigned to `label`.\"\"\"\n", + " sets = {}\n", + " for name, cfg in DATASETS.items():\n", + " target = LABELS[label][name]\n", + " rows = cfg[\"assign\"].loc[cfg[\"assign\"][\"label_name\"] == target]\n", + " assert len(rows) == PPL, f\"{name} {target}: {len(rows)} slots != {PPL}\"\n", + " sets[name] = set(rows[\"prototype_idx\"])\n", + " return sets\n", + "\n", + "\n", + "PROTOTYPES = {label: label_prototypes(label) for label in LABELS}" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## Overlap between datasets\n", + "\n", + "`region_table` splits the prototypes into the mutually exclusive regions of a Venn\n", + "diagram: `\"PTB-XL only\"` counts prototypes PTB-XL assigned to the label and no one\n", + "else, `\"PTB-XL & CinC\"` counts prototypes both of them (but not CODE-15%) assigned,\n", + "and so on." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "def region_table(sets: dict[str, set[int]]) -> pd.DataFrame:\n", + " \"\"\"Break `sets` into the mutually exclusive regions of a Venn diagram.\n", + "\n", + " Every region is listed, including empty ones, so that tables for different\n", + " labels line up. Regions are ordered by how many datasets share them.\n", + " \"\"\"\n", + " names = list(sets)\n", + " rows = []\n", + " for size in range(1, len(names) + 1):\n", + " for combo in itertools.combinations(names, size):\n", + " shared = set.intersection(*(sets[n] for n in combo))\n", + " others = [sets[n] for n in names if n not in combo]\n", + " members = shared.difference(*others) if others else shared\n", + " rows.append(\n", + " {\n", + " \"region\": \" & \".join(combo) + \" only\",\n", + " \"n_datasets\": size,\n", + " \"count\": len(members),\n", + " \"prototypes\": sorted(members),\n", + " }\n", + " )\n", + " out = pd.DataFrame(rows)\n", + " assert out[\"count\"].sum() == len(set().union(*sets.values()))\n", + " return out\n", + "\n", + "\n", + "def add_total(table: pd.DataFrame) -> pd.DataFrame:\n", + " \"\"\"Append a total row, for display only.\"\"\"\n", + " total = {\"region\": \"Total\", \"n_datasets\": None, \"count\": table[\"count\"].sum()}\n", + " return pd.concat([table, pd.DataFrame([total])], ignore_index=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "region_table(PROTOTYPES[\"Sinus Brady\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "all_brady = set()\n", + "all_tach = set()\n", + "all_af = set()\n", + "all_lbbb = set()\n", + "all_rbbb = set()\n", + "for _ds, _cfg in DATASETS.items():\n", + " _assign_df = _cfg[\"assign\"]\n", + " all_brady |= set(_assign_df.loc[(_assign_df[\"label_name\"] == LABELS[\"Sinus Brady\"][_ds]), \"prototype_idx\"])\n", + " all_tach |= set(_assign_df.loc[(_assign_df[\"label_name\"] == LABELS[\"Sinus Tach\"][_ds]), \"prototype_idx\"])\n", + " all_af |= set(_assign_df.loc[(_assign_df[\"label_name\"] == LABELS[\"AFib\"][_ds]), \"prototype_idx\"])\n", + " all_rbbb |= set(_assign_df.loc[(_assign_df[\"label_name\"] == LABELS[\"RBBB\"][_ds]), \"prototype_idx\"])\n", + " all_lbbb |= set(_assign_df.loc[(_assign_df[\"label_name\"] == LABELS[\"LBBB\"][_ds]), \"prototype_idx\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Unique AF:\", len(all_af))\n", + "print(\"Unique SB:\", len(all_brady))\n", + "print(\"Unique ST:\", len(all_tach))\n", + "\n", + "assert len(all_af & all_brady) == 0\n", + "assert len(all_af & all_tach) == 0\n", + "assert len(all_tach & all_brady) == 0\n", + "\n", + "print(\"Unique LBBB:\", len(all_lbbb))\n", + "print(\"Unique RBBB:\", len(all_rbbb))\n", + "\n", + "assert len(all_lbbb & all_rbbb) == 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Unique AF:\", len(all_af))\n", + "print(\"Unique SB:\", len(all_brady))\n", + "print(\"Unique ST:\", len(all_tach))\n", + "\n", + "assert len(all_af & all_brady) == 0\n", + "assert len(all_af & all_tach) == 0\n", + "assert len(all_tach & all_brady) == 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "# every label at once: rows are Venn regions, columns are labels\n", + "overlap = pd.DataFrame(\n", + " {\n", + " label: region_table(sets).set_index(\"region\")[\"count\"]\n", + " for label, sets in PROTOTYPES.items()\n", + " }\n", + ")\n", + "overlap.loc[\"Total Unique\"] = [\n", + " len(set().union(*sets.values())) for sets in PROTOTYPES.values()\n", + "]\n", + "overlap = overlap.rename_axis(\"Datasets\")\n", + "display(overlap)\n", + "print(overlap.to_markdown())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "# share of a label's prototypes that at least two datasets agreed on\n", + "shared = pd.DataFrame(\n", + " {\n", + " label: region_table(sets).groupby(\"n_datasets\")[\"count\"].sum()\n", + " for label, sets in PROTOTYPES.items()\n", + " }\n", + ").T\n", + "(shared[[c for c in shared.columns if c > 1]].sum(axis=1) / shared.sum(axis=1)).rename(\n", + " \"recurrent_fraction\"\n", + ").to_frame()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "# !pip install matplotlib-venn\n", + "from matplotlib.patches import Patch\n", + "from matplotlib_venn import venn3\n", + "from matplotlib_venn.layout.venn3 import DefaultLayoutAlgorithm\n", + "\n", + "assert len(DATASETS) == 3, \"venn3 draws exactly three sets\"\n", + "\n", + "TEXT_WIDTH = 5.5 # inches, width of the paper's text block\n", + "DATASET_COLORS = dict(zip(DATASETS, [\"#2a78d6\", \"#eb6834\", \"#1baf7a\"])) # colorblind-safe\n", + "\n", + "n_cols = 3\n", + "n_rows = -(-len(PROTOTYPES) // n_cols) # ceiling division\n", + "with plt.rc_context({\"font.size\": 7, \"axes.titlesize\": 8, \"axes.titleweight\": \"bold\"}):\n", + " fig, axs = plt.subplots(\n", + " n_rows, n_cols, figsize=(TEXT_WIDTH, 1.75 * n_rows + 0.3), squeeze=False\n", + " )\n", + " for ax, (label, sets) in zip(axs.flat, PROTOTYPES.items()):\n", + " # equal region sizes give every panel the same geometry\n", + " v = venn3(\n", + " [sets[name] for name in DATASETS],\n", + " set_labels=None,\n", + " set_colors=tuple(DATASET_COLORS.values()),\n", + " layout_algorithm=DefaultLayoutAlgorithm(fixed_subset_sizes=(1,) * 7),\n", + " ax=ax,\n", + " )\n", + " ax.set_title(label)\n", + " for ax in axs.flat[len(PROTOTYPES) :]: # blank out any unused panels\n", + " ax.axis(\"off\")\n", + " # take the swatches from the drawn patches so they match the translucent fill\n", + " handles = [\n", + " Patch(facecolor=v.get_patch_by_id(region).get_facecolor(), label=name)\n", + " for region, name in zip((\"100\", \"010\", \"001\"), DATASETS)\n", + " ]\n", + " fig.legend(\n", + " handles=handles,\n", + " loc=\"lower center\",\n", + " ncol=len(DATASETS),\n", + " frameon=False,\n", + " bbox_to_anchor=(0.5, 0),\n", + " )\n", + " fig.tight_layout(rect=(0, 0.06, 1, 1))\n", + " fig.savefig(FIG_DIR / f\"proto-recur-venn-seed{SEED}.png\", dpi=300, bbox_inches=\"tight\")\n", + " fig.savefig(FIG_DIR / f\"proto-recur-venn-seed{SEED}.pdf\", bbox_inches=\"tight\")" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "## Recurrent prototypes\n", + "\n", + "A prototype is recurrent for a label when at least `min_datasets` datasets picked it\n", + "for that label." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "def recurrent_prototypes(label: str, min_datasets: int = 2) -> list[int]:\n", + " \"\"\"Prototypes assigned to `label` by at least `min_datasets` datasets.\"\"\"\n", + " table = region_table(PROTOTYPES[label])\n", + " hits = table.loc[table[\"n_datasets\"] >= min_datasets, \"prototypes\"]\n", + " return sorted(p for members in hits for p in members)\n", + "\n", + "\n", + "recurrent = {label: recurrent_prototypes(label) for label in LABELS}\n", + "pd.Series(recurrent).rename(\"prototypes\").to_frame().assign(\n", + " count=lambda df: df[\"prototypes\"].str.len()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "## Plot a prototype's projections\n", + "\n", + "Each run projects a prototype onto whichever training sample and 1-second window\n", + "activates it most, so the same prototype has a different source waveform per run.\n", + "Comparing those waveforms is what shows the concept surviving across datasets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "def prototype_sources(proto_idx: int, pretrain_only: bool = False) -> pd.DataFrame:\n", + " \"\"\"Sample and window each run projected prototype `proto_idx` onto.\n", + "\n", + " Includes the HEEDB pretraining projection plus, for every dataset that assigned\n", + " the prototype, its linear probe and fine-tuned projections.\n", + " \"\"\"\n", + " rows = [{\"dataset\": \"HEEDB\", \"run\": \"pretrain\", **heedb_proj.loc[proto_idx]}]\n", + " if not pretrain_only:\n", + " for name, cfg in DATASETS.items():\n", + " slots = cfg[\"assign\"].index[cfg[\"assign\"][\"prototype_idx\"] == proto_idx]\n", + " if len(slots) == 0:\n", + " continue # this dataset did not use the prototype for any label\n", + " # ilp_effect_size uses a prototype at most once per dataset\n", + " assert len(slots) == 1, f\"{name}: prototype {proto_idx} used {len(slots)} times\"\n", + " slot = slots[0]\n", + " rows.append({\"dataset\": name, \"run\": \"probe\", **cfg[\"proj\"].loc[slot]})\n", + " rows.append({\"dataset\": name, \"run\": \"ft\", **cfg[\"ft_proj\"].loc[slot]})\n", + " columns = [\"dataset\", \"run\", \"source_id\", \"sample_id\", \"chunk_idx\", \"emb_sim\"]\n", + " ids = [\"source_id\", \"sample_id\", \"chunk_idx\"] # floated by the mixed-dtype rows\n", + " return pd.DataFrame(rows)[columns].astype(dict.fromkeys(ids, int))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "_datasets = {**DATASETS, \"HEEDB\": PRETRAIN}\n", + "_dataset_cache: dict[tuple[str, str], object] = {}\n", + "\n", + "\n", + "def get_dataset(name: str, split: str = SPLIT):\n", + " \"\"\"Build a dataset once and reuse it, loading waveforms is expensive.\"\"\"\n", + " if (name, split) not in _dataset_cache:\n", + " cfg = _datasets[name]\n", + " _dataset_cache[(name, split)] = cfg[\"dataset_cls\"](\n", + " dataset_path=str(cfg[\"dataset_path\"]),\n", + " split=split,\n", + " sampling_rate=SAMPLING_RATE,\n", + " )\n", + " return _dataset_cache[(name, split)]\n", + "\n", + "\n", + "def plot_prototype(\n", + " label: str,\n", + " proto_idx: int,\n", + " split: str = SPLIT,\n", + " save: bool = False,\n", + " pretrain_only: bool = False,\n", + ") -> dict[str, plt.Figure]:\n", + " \"\"\"Plot the source window of `proto_idx` for every run that uses it.\"\"\"\n", + " figs = {}\n", + " for _, row in prototype_sources(proto_idx, pretrain_only).iterrows():\n", + " stem = (\n", + " f\"{LABEL_SLUGS[label]}-p{proto_idx}\"\n", + " f\"-{_datasets[row['dataset']]['slug']}-{row['run']}\"\n", + " )\n", + " figs[stem] = plot_ecg(\n", + " dataset=get_dataset(row[\"dataset\"], split=split),\n", + " sample_id=int(row[\"sample_id\"]),\n", + " chunk_idx=int(row[\"chunk_idx\"]),\n", + " title=f\"{label} | prototype {proto_idx} | {row['dataset']} ({row['run']})\",\n", + " # remove_baseline=False,\n", + " )\n", + " if save:\n", + " figs[stem].savefig(FIG_DIR / f\"{stem}.png\", dpi=300, bbox_inches=\"tight\")\n", + " return figs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "# get indices of prototypes which span all datasets\n", + "for label in LABELS:\n", + " print(label, recurrent_prototypes(label, min_datasets=len(DATASETS)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": {}, + "outputs": [], + "source": [ + "# let's look at a highly-recurrent prototype\n", + "PROTO_IDX = 394\n", + "prototype_sources(PROTO_IDX)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "figs = plot_prototype(\"RBBB\", PROTO_IDX, pretrain_only=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "figs[\"rbbb-p394-heedb-pretrain\"].savefig(FIG_DIR / \"rbbb-p394-heedb-full-seed42.png\", dpi=300, bbox_inches=\"tight\")\n", + "figs[\"rbbb-p394-heedb-pretrain\"].savefig(FIG_DIR / \"rbbb-p394-heedb-full-seed42.pdf\", bbox_inches=\"tight\")" + ] + }, + { + "cell_type": "markdown", + "id": "29", + "metadata": {}, + "source": [ + "## Compact case study: projected windows only\n", + "\n", + "The full printouts repeat the same 10-second context seven times, and the underlying\n", + "ECG differs per run anyway, so for the paper we show only each run's projected\n", + "1-second window. Rows are runs (grouped by dataset) and columns are the leads that\n", + "carry the label's criteria. For RBBB those are the rsR′ in V1/V2 and the wide, slurred\n", + "S in I/V6. Each cell keeps its own y-scale.\n", + "\n", + "`has_label` checks that the source ECG carries the label. This means something only for\n", + "the HEEDB row: its projection is label-free, while the downstream projections are\n", + "label-supervised." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "from protossl.plotting import (\n", + " chunk_to_seconds,\n", + " get_ecg_waveform,\n", + " plot_segment,\n", + " resolve_index,\n", + ")\n", + "\n", + "CASE_LEADS = [\"I\", \"V1\", \"V2\", \"V6\"]\n", + "RUN_NAMES = {\"pretrain\": \"pretrained\", \"probe\": \"linear probe\", \"ft\": \"fine-tuned\"}\n", + "\n", + "\n", + "def prototype_windows(\n", + " label: str, proto_idx: int, split: str = SPLIT\n", + ") -> tuple[pd.DataFrame, list[np.ndarray]]:\n", + " \"\"\"Projected 1-second (L, T) window of `proto_idx` in every run that uses it.\n", + "\n", + " Also flags whether each window's source ECG carries `label`.\n", + " \"\"\"\n", + " sources = prototype_sources(proto_idx)\n", + " windows, has_label = [], []\n", + " for _, row in sources.iterrows():\n", + " dataset = get_dataset(row[\"dataset\"], split=split)\n", + " idx = resolve_index(dataset, sample_id=int(row[\"sample_id\"]))\n", + " signal = get_ecg_waveform(\n", + " dataset=dataset, index=idx, sampling_rate=SAMPLING_RATE\n", + " )\n", + " # same windowing as plot_ecg's chunk_idx highlight\n", + " start, end = chunk_to_seconds(\n", + " int(row[\"chunk_idx\"]),\n", + " chunk_len=SAMPLING_RATE,\n", + " chunk_overlap=0.5,\n", + " sampling_rate=SAMPLING_RATE,\n", + " )\n", + " windows.append(signal[:, int(start * SAMPLING_RATE) : int(end * SAMPLING_RATE)])\n", + " targets = list(_datasets[row[\"dataset\"]][\"targets\"])\n", + " target_idx = targets.index(LABELS[label][row[\"dataset\"]])\n", + " has_label.append(bool(dataset[idx][\"label\"][target_idx]))\n", + " return sources.assign(has_label=has_label), windows\n", + "\n", + "\n", + "def plot_prototype_windows(\n", + " sources: pd.DataFrame,\n", + " windows: list[np.ndarray],\n", + " leads: list[str] = CASE_LEADS,\n", + " row_height: float = 0.8,\n", + ") -> plt.Figure:\n", + " \"\"\"Grid of projected windows, one row per run and one column per lead.\"\"\"\n", + " # one grid row per run, plus a thin spacer row wherever the dataset changes\n", + " height_ratios, grid_rows = [], []\n", + " for i, dataset in enumerate(sources[\"dataset\"]):\n", + " if i > 0 and dataset != sources[\"dataset\"].iloc[i - 1]:\n", + " height_ratios.append(0.25)\n", + " grid_rows.append(len(height_ratios))\n", + " height_ratios.append(1)\n", + "\n", + " with plt.rc_context({\"font.size\": 7}):\n", + " fig = plt.figure(figsize=(TEXT_WIDTH, row_height * sum(height_ratios) + 0.2))\n", + " gs = fig.add_gridspec(\n", + " len(height_ratios),\n", + " len(leads),\n", + " height_ratios=height_ratios,\n", + " hspace=0.12,\n", + " wspace=0.06,\n", + " )\n", + " group_axs: dict[str, list[plt.Axes]] = {}\n", + " for (_, row), grid_row, window in zip(sources.iterrows(), grid_rows, windows):\n", + " axs = np.array([[fig.add_subplot(gs[grid_row, j]) for j in range(len(leads))]])\n", + " plot_segment(\n", + " window,\n", + " axs,\n", + " sampling_rate=SAMPLING_RATE,\n", + " lead_layout=[leads],\n", + " label_leads=False, # labeled once as column headers instead\n", + " )\n", + " axs[0, 0].set_ylabel(\n", + " RUN_NAMES[row[\"run\"]], rotation=0, ha=\"right\", va=\"center\", labelpad=4\n", + " )\n", + " group_axs.setdefault(row[\"dataset\"], []).append(axs[0, 0])\n", + " if grid_row == 0:\n", + " for ax, lead in zip(axs.flat, leads):\n", + " ax.set_title(lead, fontweight=\"bold\", fontsize=8)\n", + "\n", + " # dataset names go left of the run names, centered on each group's rows\n", + " fig.canvas.draw()\n", + " to_fig = fig.transFigure.inverted()\n", + " run_label_x0 = min(\n", + " ax.yaxis.label.get_window_extent().x0\n", + " for axs in group_axs.values()\n", + " for ax in axs\n", + " )\n", + " x = to_fig.transform((run_label_x0, 0))[0] - 0.015\n", + " for dataset, axs in group_axs.items():\n", + " y = (axs[0].get_position().y1 + axs[-1].get_position().y0) / 2\n", + " fig.text(\n", + " x, y, dataset, rotation=90, ha=\"right\", va=\"center\",\n", + " fontweight=\"bold\", fontsize=8,\n", + " )\n", + " return fig\n", + "\n", + "\n", + "sources, windows = prototype_windows(\"RBBB\", PROTO_IDX)\n", + "display(sources)\n", + "assert sources[\"has_label\"].all()\n", + "\n", + "fig = plot_prototype_windows(sources, windows)\n", + "stem = f\"{LABEL_SLUGS['RBBB']}-p{PROTO_IDX}-windows-seed{SEED}\"\n", + "fig.savefig(FIG_DIR / f\"{stem}.png\", dpi=300, bbox_inches=\"tight\")\n", + "fig.savefig(FIG_DIR / f\"{stem}.pdf\", bbox_inches=\"tight\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "protossl", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/scripts/2-run-labsup-proto-direct.sh b/scripts/2-run-supproto-direct.sh similarity index 98% rename from scripts/2-run-labsup-proto-direct.sh rename to scripts/2-run-supproto-direct.sh index cd928b1..7b4f97c 100755 --- a/scripts/2-run-labsup-proto-direct.sh +++ b/scripts/2-run-supproto-direct.sh @@ -14,7 +14,7 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="labsup-proto-direct" +EXP_NAME="supproto-direct" python -m protossl.trainer \ --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ diff --git a/scripts/2-z-run-labsup-proto-direct-28ppl.sh b/scripts/2-z-run-supproto-direct-28ppl.sh similarity index 98% rename from scripts/2-z-run-labsup-proto-direct-28ppl.sh rename to scripts/2-z-run-supproto-direct-28ppl.sh index 6e84c18..b52ee09 100755 --- a/scripts/2-z-run-labsup-proto-direct-28ppl.sh +++ b/scripts/2-z-run-supproto-direct-28ppl.sh @@ -14,7 +14,7 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="labsup-proto-direct-28ppl" +EXP_NAME="supproto-direct-28ppl" if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/2-z-run-labsup-proto-direct-7ppl.sh b/scripts/2-z-run-supproto-direct-7ppl.sh similarity index 98% rename from scripts/2-z-run-labsup-proto-direct-7ppl.sh rename to scripts/2-z-run-supproto-direct-7ppl.sh index 8a9fa59..a6e88bd 100755 --- a/scripts/2-z-run-labsup-proto-direct-7ppl.sh +++ b/scripts/2-z-run-supproto-direct-7ppl.sh @@ -14,7 +14,7 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="labsup-proto-direct-7ppl" +EXP_NAME="supproto-direct-7ppl" if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/3-run-protossl-heedb-150-1050-pila.sh b/scripts/3-run-protossl-heedb-150-1050-pila.sh new file mode 100755 index 0000000..6bcf511 --- /dev/null +++ b/scripts/3-run-protossl-heedb-150-1050-pila.sh @@ -0,0 +1,110 @@ +#!/bin/bash + +set -e + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# experiment parameters +EXP_NAME="protossl-heedb-150-1050-pila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb-150-1050") + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --model.n_prototypes 1050 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototypes/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +# now fine-tune +PRETRAIN_RUN=$RUN_DIR/$EXP_NAME +EXP_NAME="$EXP_NAME-ft" + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/3-run-protossl-heedb-150-pila.sh b/scripts/3-run-protossl-heedb-150-pila.sh new file mode 100755 index 0000000..372e1f6 --- /dev/null +++ b/scripts/3-run-protossl-heedb-150-pila.sh @@ -0,0 +1,110 @@ +#!/bin/bash + +set -e + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# experiment parameters +EXP_NAME="protossl-heedb-150-pila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb-150") + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --model.n_prototypes 7500 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototypes/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +# now fine-tune +PRETRAIN_RUN=$RUN_DIR/$EXP_NAME +EXP_NAME="$EXP_NAME-ft" + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/3-run-protossl-heedb-pila.sh b/scripts/3-run-protossl-heedb-pila.sh index feaaa71..923ffd6 100755 --- a/scripts/3-run-protossl-heedb-pila.sh +++ b/scripts/3-run-protossl-heedb-pila.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pila" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") python -m protossl.trainer \ --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ diff --git a/scripts/3-z-run-protossl-heedb-pila-28ppl.sh b/scripts/3-z-run-protossl-heedb-pila-28ppl.sh index c780971..89d48c5 100755 --- a/scripts/3-z-run-protossl-heedb-pila-28ppl.sh +++ b/scripts/3-z-run-protossl-heedb-pila-28ppl.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pila-28ppl" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/3-z-run-protossl-heedb-pila-7ppl.sh b/scripts/3-z-run-protossl-heedb-pila-7ppl.sh index 149586d..41304fd 100755 --- a/scripts/3-z-run-protossl-heedb-pila-7ppl.sh +++ b/scripts/3-z-run-protossl-heedb-pila-7ppl.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pila-7ppl" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/3-z-run-protossl-heedb-pila-83ppl.sh b/scripts/3-z-run-protossl-heedb-pila-83ppl.sh index 876ca9a..a823f16 100755 --- a/scripts/3-z-run-protossl-heedb-pila-83ppl.sh +++ b/scripts/3-z-run-protossl-heedb-pila-83ppl.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pila-83ppl" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/4-run-supproto-heedb-150-1050-rila.sh b/scripts/4-run-supproto-heedb-150-1050-rila.sh new file mode 100755 index 0000000..a9eff48 --- /dev/null +++ b/scripts/4-run-supproto-heedb-150-1050-rila.sh @@ -0,0 +1,118 @@ +#!/bin/bash + +set -e + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# experiment parameters +EXP_NAME="supproto-heedb-150-1050-rila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-heedb-150-1050") + +if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --model.n_prototypes 1050 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/project-prototypes-supervised/latest/proj.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + + python _linear_probe.py \ + --random-seed $SEED \ + --dataset-path $DATASET_PATH \ + --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs_bootstrapped.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + +fi + +# now fine-tune +PRETRAIN_RUN=$RUN_DIR/$EXP_NAME +EXP_NAME="$EXP_NAME-ft" + +if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + + python _linear_probe.py \ + --random-seed $SEED \ + --dataset-path $DATASET_PATH \ + --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs_bootstrapped.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + +fi diff --git a/scripts/4-run-supproto-heedb-150-rila.sh b/scripts/4-run-supproto-heedb-150-rila.sh new file mode 100755 index 0000000..be934c4 --- /dev/null +++ b/scripts/4-run-supproto-heedb-150-rila.sh @@ -0,0 +1,118 @@ +#!/bin/bash + +set -e + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# experiment parameters +EXP_NAME="supproto-heedb-150-rila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-heedb-150") + +if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --model.n_prototypes 7500 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/project-prototypes-supervised/latest/proj.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + + python _linear_probe.py \ + --random-seed $SEED \ + --dataset-path $DATASET_PATH \ + --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs_bootstrapped.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + +fi + +# now fine-tune +PRETRAIN_RUN=$RUN_DIR/$EXP_NAME +EXP_NAME="$EXP_NAME-ft" + +if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + + python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + + python _linear_probe.py \ + --random-seed $SEED \ + --dataset-path $DATASET_PATH \ + --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + + python _eval_probs_bootstrapped.py \ + --dataset-path $DATASET_PATH \ + --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ + --output-path $RUN_DIR/$EXP_NAME + +fi diff --git a/scripts/4-run-labsup-proto-heedb-rila.sh b/scripts/4-run-supproto-heedb-rila.sh similarity index 97% rename from scripts/4-run-labsup-proto-heedb-rila.sh rename to scripts/4-run-supproto-heedb-rila.sh index 3c49364..1b30aac 100755 --- a/scripts/4-run-labsup-proto-heedb-rila.sh +++ b/scripts/4-run-supproto-heedb-rila.sh @@ -14,8 +14,8 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="labsup-proto-heedb-rila" -PRETRAIN_RUN="$RUN_DIR/../prosup-pretrain-heedb" +EXP_NAME="supproto-heedb-rila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/4-y-run-labsup-proto-heedb-pila.sh b/scripts/4-y-run-supproto-heedb-pila.sh similarity index 97% rename from scripts/4-y-run-labsup-proto-heedb-pila.sh rename to scripts/4-y-run-supproto-heedb-pila.sh index 2a9a70d..9b9b1e1 100755 --- a/scripts/4-y-run-labsup-proto-heedb-pila.sh +++ b/scripts/4-y-run-supproto-heedb-pila.sh @@ -14,8 +14,8 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="labsup-proto-heedb-pila" -PRETRAIN_RUN="$RUN_DIR/../prosup-pretrain-heedb" +EXP_NAME="supproto-heedb-pila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/4-z-run-labsup-proto-heedb-rila-28ppl.sh b/scripts/4-z-run-supproto-heedb-rila-28ppl.sh similarity index 97% rename from scripts/4-z-run-labsup-proto-heedb-rila-28ppl.sh rename to scripts/4-z-run-supproto-heedb-rila-28ppl.sh index d2ba9a6..18047a4 100755 --- a/scripts/4-z-run-labsup-proto-heedb-rila-28ppl.sh +++ b/scripts/4-z-run-supproto-heedb-rila-28ppl.sh @@ -14,8 +14,8 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="labsup-proto-heedb-rila-28ppl" -PRETRAIN_RUN="$RUN_DIR/../prosup-pretrain-heedb" +EXP_NAME="supproto-heedb-rila-28ppl" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/4-z-run-labsup-proto-heedb-rila-7ppl.sh b/scripts/4-z-run-supproto-heedb-rila-7ppl.sh similarity index 97% rename from scripts/4-z-run-labsup-proto-heedb-rila-7ppl.sh rename to scripts/4-z-run-supproto-heedb-rila-7ppl.sh index c686fac..00fddee 100755 --- a/scripts/4-z-run-labsup-proto-heedb-rila-7ppl.sh +++ b/scripts/4-z-run-supproto-heedb-rila-7ppl.sh @@ -14,8 +14,8 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="labsup-proto-heedb-rila-7ppl" -PRETRAIN_RUN="$RUN_DIR/../prosup-pretrain-heedb" +EXP_NAME="supproto-heedb-rila-7ppl" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/5-3-run-pclr-logreg.sh b/scripts/5-3-run-pclr-logreg.sh new file mode 100755 index 0000000..4879251 --- /dev/null +++ b/scripts/5-3-run-pclr-logreg.sh @@ -0,0 +1,39 @@ +#!/bin/bash + +set -e + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# experiment parameters +EXP_NAME="pclr-logreg" +PRETRAIN_RUN=$(realpath -m $RUN_DIR/../../pclr-heedb) + +python ecg-fms/_compute_pclr_embeddings.py \ +--dataset-path $DATASET_PATH \ +--ckpt $PRETRAIN_RUN/train-contraster/latest/best.ckpt \ +--output-path $RUN_DIR/$EXP_NAME + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/6-run-protossl-heedb-pia.sh b/scripts/6-run-protossl-heedb-pia.sh index a067e3e..60ae22e 100755 --- a/scripts/6-run-protossl-heedb-pia.sh +++ b/scripts/6-run-protossl-heedb-pia.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pia" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") python -m protossl.trainer \ --config $REPO_ROOT/configs/target-guided-14ppl.yaml \ diff --git a/scripts/6-z-run-protossl-heedb-pia-83ppl.sh b/scripts/6-z-run-protossl-heedb-pia-83ppl.sh index a231c00..1ab9383 100755 --- a/scripts/6-z-run-protossl-heedb-pia-83ppl.sh +++ b/scripts/6-z-run-protossl-heedb-pia-83ppl.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pia-83ppl" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") if [ ! -e "$RUN_DIR/$EXP_NAME/metrics-bootstrapped.csv" ]; then diff --git a/scripts/7-run-protossl-heedb-pit.sh b/scripts/7-1-run-protossl-heedb-pit.sh similarity index 96% rename from scripts/7-run-protossl-heedb-pit.sh rename to scripts/7-1-run-protossl-heedb-pit.sh index 0035093..acefad5 100755 --- a/scripts/7-run-protossl-heedb-pit.sh +++ b/scripts/7-1-run-protossl-heedb-pit.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pit" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") python -m protossl.trainer \ --config $REPO_ROOT/configs/target-unguided.yaml \ diff --git a/scripts/8-run-protossl-heedb-pip.sh b/scripts/7-2-run-protossl-heedb-pip.sh similarity index 95% rename from scripts/8-run-protossl-heedb-pip.sh rename to scripts/7-2-run-protossl-heedb-pip.sh index ea25fee..a9e84c8 100755 --- a/scripts/8-run-protossl-heedb-pip.sh +++ b/scripts/7-2-run-protossl-heedb-pip.sh @@ -15,7 +15,7 @@ cd $REPO_ROOT/scripts # experiment parameters EXP_NAME="protossl-heedb-pip" -PRETRAIN_RUN="$RUN_DIR/../pass-pretrain-heedb-no-attn" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") python -m protossl.trainer \ --config $REPO_ROOT/configs/target-unguided.yaml \ diff --git a/scripts/7-3-run-protossl-heedb-latent.sh b/scripts/7-3-run-protossl-heedb-latent.sh new file mode 100755 index 0000000..0c00432 --- /dev/null +++ b/scripts/7-3-run-protossl-heedb-latent.sh @@ -0,0 +1,43 @@ +#!/bin/bash + +set -e + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# experiment parameters +EXP_NAME="protossl-heedb-latent" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-heedb") + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/target-unguided.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototypes/latest/best.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/README.md b/scripts/README.md index 6af1550..30a7b1c 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -7,6 +7,6 @@ We rarely call individual run scripts directly (with one exception, detailed bel The only run scripts we execute directly (or submit manually to slurm) are the pretraining jobs under the `pretrain` subdirectory. If you're trying to replicate our experiments, we'd suggest the following: -1. Find and replace paths which start with `/opt/gpudata/steven` - these are the only places which contain hardcoded paths and should only be present within a handful of run scripts -1. Run `pretrain/run-pass-heedb-pretrain.sh` +1. Find and replace paths which start with `/opt/gpu_working/steven` (repo and output directories) and `/opt/gpudata` (source datasets) - these are the only places which contain hardcoded paths and should only be present within the run, queue, and pretrain scripts +1. Run `pretrain/run-protossl-heedb-pretrain.sh` 1. Run all `queue-*-experiments.sh` diff --git a/scripts/_submit_job.sh b/scripts/_submit_job.sh index 2feb4d2..5701550 100644 --- a/scripts/_submit_job.sh +++ b/scripts/_submit_job.sh @@ -18,7 +18,7 @@ function submit_job() { local fname="${2%.sh}" - local job_id=$(sbatch -p h200 --gpus=1 \ + local job_id=$(sbatch -w kg35-nvl01 --gpus=1 \ --time=0 --cpus-per-task=24 --mem=200g --ntasks=1 \ --export=ALL,DATASET_PATH=$BASE_DATASET_PATH$1,RUN_DIR=$BASE_RUN_DIR$1,REPO_ROOT=$REPO_ROOT,SEED=$SEED \ --output $HOME/slurm-logs/%j-$fname$1-%N.out $3 --parsable _slurm_wrapper.sh $REPO_ROOT/scripts/$2) diff --git a/scripts/audio/2-run-labsup-proto-direct.sh b/scripts/audio/2-run-supproto-direct.sh similarity index 98% rename from scripts/audio/2-run-labsup-proto-direct.sh rename to scripts/audio/2-run-supproto-direct.sh index bc6420b..979a71f 100755 --- a/scripts/audio/2-run-labsup-proto-direct.sh +++ b/scripts/audio/2-run-supproto-direct.sh @@ -27,7 +27,7 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts/audio # experiment parameters -EXP_NAME="labsup-proto-direct" +EXP_NAME="supproto-direct" # python -m protossl.trainer \ # --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ diff --git a/scripts/audio/3-run-protossl-audioset-pila.sh b/scripts/audio/3-run-protossl-audioset-pila.sh index eb0e8c8..0fd1a13 100755 --- a/scripts/audio/3-run-protossl-audioset-pila.sh +++ b/scripts/audio/3-run-protossl-audioset-pila.sh @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=200gb +#SBATCH --mem-per-gpu=100gb #SBATCH --gpus-per-node=1 #SBATCH --nodes=1 #SBATCH -w kg35-nvl01 @@ -28,48 +28,48 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="protossl-audioset-pila" -PRETRAIN_RUN="$RUN_DIR/../pass-audioset" - -# python -m protossl.trainer \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototype-assignments \ -# --assignment-strategy ilp_effect_size \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --model.n_prototypes 2635 \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/learn-prototypes/latest/best.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass"}' - -# python -m protossl.trainer \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt - -# python -m protossl.trainer \ -# --seed_everything $SEED \ -# --pipeline-stage compute-embeddings \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# python _linear_probe.py \ -# --random-seed $SEED \ -# --dataset-path $DATASET_PATH \ -# --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ -# --output-path $RUN_DIR/$EXP_NAME - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2") + +python -m protossl.trainer \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --model.n_prototypes 6324 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototypes/latest/best.ckpt \ + --model.model_kwargs '{"label_type": "multiclass"}' + +python -m protossl.trainer \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME python _eval_probs_bootstrapped.py \ --dataset-path $DATASET_PATH \ @@ -80,44 +80,44 @@ python _eval_probs_bootstrapped.py \ PRETRAIN_RUN=$RUN_DIR/$EXP_NAME EXP_NAME="$EXP_NAME-ft" -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage compute-embeddings \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# python _linear_probe.py \ -# --random-seed $SEED \ -# --dataset-path $DATASET_PATH \ -# --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ -# --output-path $RUN_DIR/$EXP_NAME - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ + --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME python _eval_probs_bootstrapped.py \ --dataset-path $DATASET_PATH \ diff --git a/scripts/audio/4-run-labsup-proto-audioset-rila.sh b/scripts/audio/4-run-labsup-proto-audioset-rila.sh deleted file mode 100755 index a696997..0000000 --- a/scripts/audio/4-run-labsup-proto-audioset-rila.sh +++ /dev/null @@ -1,125 +0,0 @@ -#!/bin/bash - -#SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=200gb -#SBATCH --gpus-per-node=1 -#SBATCH --nodes=1 -#SBATCH -w kg35-nvl01 -#SBATCH --ntasks-per-node=1 -#SBATCH --time=0 -#SBATCH --output /home/songs1/slurm-logs/audio-%j.out - -set -e - -export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL - -# set these env vars prior to executing this script -: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" -: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" -: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" -: "${SEED:=42}" -: "${PPL:=5ppl}" -echo "Using PPL=$PPL" -echo "Using SEED=$SEED" -echo "Using DATASET_PATH=$DATASET_PATH" -echo "Using RUN_DIR=$RUN_DIR" -echo "Using REPO_ROOT=$REPO_ROOT" -cd $REPO_ROOT/scripts/audio - -# experiment parameters -EXP_NAME="labsup-proto-audioset-rila" -PRETRAIN_RUN="$RUN_DIR/../prosup-audioset" - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototype-assignments \ -# --assignment-strategy ilp_effect_size \ -# --model.n_prototypes 2635 \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/project-prototypes-supervised/latest/proj.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass"}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage compute-embeddings \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# python _linear_probe.py \ -# --random-seed $SEED \ -# --dataset-path $DATASET_PATH \ -# --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ -# --output-path $RUN_DIR/$EXP_NAME - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME - -python _eval_probs_bootstrapped.py \ ---dataset-path $DATASET_PATH \ ---probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ ---output-path $RUN_DIR/$EXP_NAME - -# now fine-tune -PRETRAIN_RUN=$RUN_DIR/$EXP_NAME -EXP_NAME="$EXP_NAME-ft" - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage compute-embeddings \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# python _linear_probe.py \ -# --random-seed $SEED \ -# --dataset-path $DATASET_PATH \ -# --prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ -# --output-path $RUN_DIR/$EXP_NAME - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME - -python _eval_probs_bootstrapped.py \ ---dataset-path $DATASET_PATH \ ---probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ ---output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/audio/4-run-supproto-audioset-rila.sh b/scripts/audio/4-run-supproto-audioset-rila.sh new file mode 100755 index 0000000..3b01623 --- /dev/null +++ b/scripts/audio/4-run-supproto-audioset-rila.sh @@ -0,0 +1,125 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=100gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/audio-%j.out + +set -e + +export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +: "${PPL:=5ppl}" +echo "Using PPL=$PPL" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts/audio + +# experiment parameters +EXP_NAME="supproto-audioset-rila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2") + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --model.n_prototypes 6324 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/project-prototypes-supervised/latest/proj.ckpt \ + --model.model_kwargs '{"label_type": "multiclass"}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +# now fine-tune +PRETRAIN_RUN=$RUN_DIR/$EXP_NAME +EXP_NAME="$EXP_NAME-ft" + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ + --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage compute-embeddings \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +python _linear_probe.py \ +--random-seed $SEED \ +--dataset-path $DATASET_PATH \ +--prototype-embeddings $RUN_DIR/$EXP_NAME/compute-embeddings/latest \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/audio/queue-experiments.sh b/scripts/audio/queue-experiments.sh index ff3acb4..8f65d2a 100644 --- a/scripts/audio/queue-experiments.sh +++ b/scripts/audio/queue-experiments.sh @@ -2,23 +2,23 @@ scripts=( "1-run-blackbox-direct.sh" - "2-run-labsup-proto-direct.sh" + "2-run-supproto-direct.sh" "3-run-protossl-audioset-pila.sh" - "4-run-labsup-proto-audioset-rila.sh" + "4-run-supproto-audioset-rila.sh" ) voxceleb_scripts=( "1-run-blackbox-direct.sh" - "voxceleb-2-run-labsup-proto-direct.sh" + "voxceleb-2-run-supproto-direct.sh" "voxceleb-3-run-protossl-audioset-pila.sh" - "voxceleb-4-run-labsup-proto-audioset-rila.sh" + "voxceleb-4-run-supproto-audioset-rila.sh" ) # iemocap for i in {0..4}; do export IEMOCAP_TEST_FOLD=$i export DATASET_PATH=/opt/gpudata/audio/IEMOCAP - export RUN_DIR=/opt/gpu_working/steven/protossl-audio/runs-iemocap-fold$i + export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-iemocap-fold$i-v2 for script in "${scripts[@]}"; do sbatch $script done @@ -28,7 +28,7 @@ done for i in {0..9}; do export US8K_TEST_FOLD=$i export DATASET_PATH=/opt/gpudata/audio/UrbanSound8K - export RUN_DIR=/opt/gpu_working/steven/protossl-audio/runs-us8k-fold$i + export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-us8k-fold$i-v2 for script in "${scripts[@]}"; do sbatch $script done @@ -36,8 +36,8 @@ done # voxceleb1 (these take even longer to run) export DATASET_PATH=/opt/gpudata/audio/VoxCeleb1 -export RUN_DIR=/opt/gpu_working/steven/protossl-audio/runs-voxceleb -export PPL=2ppl +export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-voxceleb-v2 +export PPL=5ppl for script in "${voxceleb_scripts[@]}"; do sbatch $script done diff --git a/scripts/audio/run-pretrain-audio-pass.sh b/scripts/audio/run-pretrain-audio-protossl.sh similarity index 77% rename from scripts/audio/run-pretrain-audio-pass.sh rename to scripts/audio/run-pretrain-audio-protossl.sh index 824001c..f88908b 100644 --- a/scripts/audio/run-pretrain-audio-pass.sh +++ b/scripts/audio/run-pretrain-audio-protossl.sh @@ -12,7 +12,7 @@ set -e export DATASET_PATH=/opt/gpudata/audio/audioset -export RUN_DIR=/opt/gpu_working/steven/protossl-audio +export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL : "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" @@ -26,7 +26,7 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="pass-audioset" +EXP_NAME="protossl-audioset-v2" echo "======================================" echo "Stage 1: self-supervised learn-prototypes" @@ -34,7 +34,8 @@ echo "======================================" srun python -m protossl.trainer \ --pipeline-stage learn-prototypes \ - --config $REPO_ROOT/configs/audio/pretrain-unsupervised.yaml \ + --config $REPO_ROOT/configs/audio/pretrain-unsupervised-v2.yaml \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ - --data.dataset_path $DATASET_PATH + --data.dataset_path $DATASET_PATH \ + --resume-from-checkpoint $RUN_DIR/$EXP_NAME/learn-prototypes/v2-s73uvkdj/best.ckpt diff --git a/scripts/audio/run-pretrain-audio-prosup.sh b/scripts/audio/run-pretrain-audio-supproto.sh similarity index 69% rename from scripts/audio/run-pretrain-audio-prosup.sh rename to scripts/audio/run-pretrain-audio-supproto.sh index 96c62ac..fde1e22 100644 --- a/scripts/audio/run-pretrain-audio-prosup.sh +++ b/scripts/audio/run-pretrain-audio-supproto.sh @@ -12,7 +12,7 @@ set -e export DATASET_PATH=/opt/gpudata/audio/audioset -export RUN_DIR=/opt/gpu_working/steven/protossl-audio +export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL : "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" @@ -26,19 +26,20 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts # experiment parameters -EXP_NAME="prosup-audioset" +EXP_NAME="supproto-audioset-v2" echo "======================================" echo "Stage 1: learn-prototypes-supervised" echo "======================================" -# srun python -m protossl.trainer \ -# --pipeline-stage learn-prototypes-supervised \ -# --config $REPO_ROOT/configs/audio/pretrain-supervised.yaml \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.model_kwargs '{"use_default_weights": True}' +srun python -m protossl.trainer \ + --pipeline-stage learn-prototypes-supervised \ + --config $REPO_ROOT/configs/audio/pretrain-supervised-v2.yaml \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.model_kwargs '{"use_default_weights": True}' \ + --resume-from-checkpoint $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/v1-iw8u33te/best.ckpt echo "======================================" echo "Stage 2: project prototypes supervised" @@ -46,7 +47,7 @@ echo "======================================" srun python -m protossl.trainer \ --pipeline-stage project-prototypes-supervised \ - --config $REPO_ROOT/configs/audio/pretrain-supervised.yaml \ + --config $REPO_ROOT/configs/audio/pretrain-supervised-v2.yaml \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ diff --git a/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh b/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh deleted file mode 100755 index fc716a0..0000000 --- a/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh +++ /dev/null @@ -1,74 +0,0 @@ -#!/bin/bash - -#SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=200gb -#SBATCH --gpus-per-node=1 -#SBATCH --nodes=1 -#SBATCH -w kg35-nvl01 -#SBATCH --ntasks-per-node=1 -#SBATCH --time=0 -#SBATCH --output /home/songs1/slurm-logs/audio-%j.out - -set -e - -export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL - -export DATASET_PATH=/opt/gpudata/audio/VoxCeleb1 -export RUN_DIR=/opt/gpu_working/steven/protossl-audio/runs-voxceleb -export PPL=2ppl - -# set these env vars prior to executing this script -: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" -: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" -: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" -: "${SEED:=42}" -: "${PPL:=5ppl}" -echo "Using PPL=$PPL" -echo "Using SEED=$SEED" -echo "Using DATASET_PATH=$DATASET_PATH" -echo "Using RUN_DIR=$RUN_DIR" -echo "Using REPO_ROOT=$REPO_ROOT" -cd $REPO_ROOT/scripts/audio - -# experiment parameters -EXP_NAME="labsup-proto-direct" - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage train-classifier \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.model_kwargs '{"label_type": "multiclass"}' \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME - -python _eval_probs_bootstrapped.py \ ---dataset-path $DATASET_PATH \ ---probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ ---output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/audio/voxceleb-2-run-supproto-direct.sh b/scripts/audio/voxceleb-2-run-supproto-direct.sh new file mode 100755 index 0000000..272973a --- /dev/null +++ b/scripts/audio/voxceleb-2-run-supproto-direct.sh @@ -0,0 +1,74 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=100gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/audio-%j.out + +set -e + +export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +export DATASET_PATH=/opt/gpudata/audio/VoxCeleb1 +export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-voxceleb-v2 +export PPL=5ppl + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +: "${PPL:=5ppl}" +echo "Using PPL=$PPL" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts/audio + +# experiment parameters +EXP_NAME="supproto-direct" + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage train-classifier \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.model_kwargs '{"label_type": "multiclass"}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh b/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh index 7158ae9..b3c9d71 100755 --- a/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh +++ b/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=200gb +#SBATCH --mem-per-gpu=100gb #SBATCH --gpus-per-node=1 #SBATCH --nodes=1 #SBATCH -w kg35-nvl01 @@ -14,8 +14,8 @@ set -e export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL export DATASET_PATH=/opt/gpudata/audio/VoxCeleb1 -export RUN_DIR=/opt/gpu_working/steven/protossl-audio/runs-voxceleb -export PPL=2ppl +export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-voxceleb-v2 +export PPL=5ppl # set these env vars prior to executing this script : "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" @@ -32,45 +32,45 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="protossl-audioset-pila" -PRETRAIN_RUN="$RUN_DIR/../pass-audioset" - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototype-assignments \ -# --assignment-strategy ilp_effect_size \ -# --model.n_prototypes 2635 \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/learn-prototypes/latest/best.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass"}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage train-classifier \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.model_kwargs '{"label_type": "multiclass"}' \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2") + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --model.n_prototypes 6324 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototypes/latest/best.ckpt \ + --model.model_kwargs '{"label_type": "multiclass"}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage train-classifier \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.model_kwargs '{"label_type": "multiclass"}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME python _eval_probs_bootstrapped.py \ --dataset-path $DATASET_PATH \ @@ -81,41 +81,41 @@ python _eval_probs_bootstrapped.py \ PRETRAIN_RUN=$RUN_DIR/$EXP_NAME EXP_NAME="$EXP_NAME-ft" -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage train-classifier \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.model_kwargs '{"label_type": "multiclass"}' \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ + --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage train-classifier \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.model_kwargs '{"label_type": "multiclass"}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME python _eval_probs_bootstrapped.py \ --dataset-path $DATASET_PATH \ diff --git a/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh b/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh deleted file mode 100755 index 5fd8ab9..0000000 --- a/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh +++ /dev/null @@ -1,123 +0,0 @@ -#!/bin/bash - -#SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=200gb -#SBATCH --gpus-per-node=1 -#SBATCH --nodes=1 -#SBATCH -w kg35-nvl01 -#SBATCH --ntasks-per-node=1 -#SBATCH --time=0 -#SBATCH --output /home/songs1/slurm-logs/audio-%j.out - -set -e - -export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL - -export DATASET_PATH=/opt/gpudata/audio/VoxCeleb1 -export RUN_DIR=/opt/gpu_working/steven/protossl-audio/runs-voxceleb -export PPL=2ppl - -# set these env vars prior to executing this script -: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" -: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" -: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" -: "${SEED:=42}" -: "${PPL:=5ppl}" -echo "Using PPL=$PPL" -echo "Using SEED=$SEED" -echo "Using DATASET_PATH=$DATASET_PATH" -echo "Using RUN_DIR=$RUN_DIR" -echo "Using REPO_ROOT=$REPO_ROOT" -cd $REPO_ROOT/scripts/audio - -# experiment parameters -EXP_NAME="labsup-proto-audioset-rila" -PRETRAIN_RUN="$RUN_DIR/../prosup-audioset" - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototype-assignments \ -# --assignment-strategy ilp_effect_size \ -# --model.n_prototypes 2635 \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/project-prototypes-supervised/latest/proj.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass"}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage train-classifier \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.model_kwargs '{"label_type": "multiclass"}' \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME - -python _eval_probs_bootstrapped.py \ ---dataset-path $DATASET_PATH \ ---probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ ---output-path $RUN_DIR/$EXP_NAME - -# now fine-tune -PRETRAIN_RUN=$RUN_DIR/$EXP_NAME -EXP_NAME="$EXP_NAME-ft" - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage learn-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ -# --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage project-prototypes-supervised \ -# --trainer.logger.save_dir $RUN_DIR/ \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt - -# python -m protossl.trainer \ -# --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ -# --seed_everything $SEED \ -# --pipeline-stage train-classifier \ -# --trainer.logger.save_dir $RUN_DIR \ -# --trainer.logger.name $EXP_NAME \ -# --data.dataset_path $DATASET_PATH \ -# --model.model_kwargs '{"label_type": "multiclass"}' \ -# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt - -# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy - -# python _eval_probs.py \ -# --dataset-path $DATASET_PATH \ -# --probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ -# --output-path $RUN_DIR/$EXP_NAME - -python _eval_probs_bootstrapped.py \ ---dataset-path $DATASET_PATH \ ---probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ ---output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/audio/voxceleb-4-run-supproto-audioset-rila.sh b/scripts/audio/voxceleb-4-run-supproto-audioset-rila.sh new file mode 100755 index 0000000..77e7e88 --- /dev/null +++ b/scripts/audio/voxceleb-4-run-supproto-audioset-rila.sh @@ -0,0 +1,123 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=100gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/audio-%j.out + +set -e + +export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +export DATASET_PATH=/opt/gpudata/audio/VoxCeleb1 +export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-voxceleb-v2 +export PPL=5ppl + +# set these env vars prior to executing this script +: "${DATASET_PATH:?Env var DATASET_PATH must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +: "${SEED:=42}" +: "${PPL:=5ppl}" +echo "Using PPL=$PPL" +echo "Using SEED=$SEED" +echo "Using DATASET_PATH=$DATASET_PATH" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts/audio + +# experiment parameters +EXP_NAME="supproto-audioset-rila" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2") + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototype-assignments \ + --assignment-strategy ilp_effect_size \ + --model.n_prototypes 6324 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/project-prototypes-supervised/latest/proj.ckpt \ + --model.model_kwargs '{"label_type": "multiclass"}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototype-assignments/latest/assigned.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage train-classifier \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.model_kwargs '{"label_type": "multiclass"}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +# now fine-tune +PRETRAIN_RUN=$RUN_DIR/$EXP_NAME +EXP_NAME="$EXP_NAME-ft" + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage learn-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $PRETRAIN_RUN/learn-prototype-assignments/latest/assigned.ckpt \ + --model.model_kwargs '{"label_type": "multiclass", "use_default_weights": True}' + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage project-prototypes-supervised \ + --trainer.logger.save_dir $RUN_DIR/ \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt + +python -m protossl.trainer \ + --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ + --seed_everything $SEED \ + --pipeline-stage train-classifier \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $DATASET_PATH \ + --model.model_kwargs '{"label_type": "multiclass"}' \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy + +python _eval_probs.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--probs-npy $RUN_DIR/$EXP_NAME/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME diff --git a/scripts/ecg-fms/_compute_ecgfounder_embeddings.py b/scripts/ecg-fms/_compute_ecgfounder_embeddings.py index d138ace..40de11c 100644 --- a/scripts/ecg-fms/_compute_ecgfounder_embeddings.py +++ b/scripts/ecg-fms/_compute_ecgfounder_embeddings.py @@ -56,7 +56,7 @@ def main( # based on dataset implementation in ECGFounder/ptbxl_eval.py # the only transformation is normalization and resampling - # which we do in the pass_pclr dataset implementations + # which we do in the protossl dataset implementations ds_cls, _, is_audio = infer_dataset_class_from_path(dataset_path) assert not is_audio train_ds = ds_cls( diff --git a/scripts/ecg-fms/_compute_pclr_embeddings.py b/scripts/ecg-fms/_compute_pclr_embeddings.py new file mode 100644 index 0000000..97c73fc --- /dev/null +++ b/scripts/ecg-fms/_compute_pclr_embeddings.py @@ -0,0 +1,87 @@ +import argparse +import os + +import numpy as np +import torch +from torch.utils.data import DataLoader +from tqdm import tqdm + +from protossl.datasets import infer_dataset_class_from_path +from protossl.models import BlackboxContraster + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--dataset-path", required=True) + parser.add_argument("--output-path", required=True) + parser.add_argument("--ckpt", required=True) + parser.add_argument("--batch-size", type=int, default=512) + args = parser.parse_args() + return args + + +def main( + *, # enforce kwargs + dataset_path: str, + output_path: str, + ckpt_path: str, + batch_size: int = 512, +): + # load ckpt for hyperparameters + ckpt = torch.load(ckpt_path, map_location="cpu") + hparams: dict = ckpt["hyper_parameters"] # type: ignore + data_hparams: dict = ckpt["datamodule_hyper_parameters"] # type: ignore + + # load and freeze model + model = BlackboxContraster( + contrastive_pair_mode=hparams["contrastive_pair_mode"], + backbone_type=hparams["backbone_type"], + conv_type=hparams["conv_type"], + **hparams["model_kwargs"], + pretrained_weights=ckpt_path, # internally loads checkpoint again, but ok as model is small + ) + for param in model.parameters(): + param.requires_grad = False + + device = "cuda" if torch.cuda.is_available() else "cpu" + encoder = model.encoder + encoder = encoder.eval().to(device) + + ds_cls, _, is_audio = infer_dataset_class_from_path(dataset_path) + assert not is_audio + train_ds = ds_cls( + dataset_path=dataset_path, + split="train", + sampling_rate=data_hparams["sampling_rate"], + ) + test_ds = ds_cls( + dataset_path=dataset_path, + split="test", + sampling_rate=data_hparams["sampling_rate"], + ) + train_dl = DataLoader(train_ds, batch_size=batch_size) + test_dl = DataLoader(test_ds, batch_size=batch_size) + + train_embs, test_embs = [], [] + with torch.inference_mode(): + for dl, embs in [(train_dl, train_embs), (test_dl, test_embs)]: + for batch in tqdm(dl): + x = batch["waveform"].to(device) + batch_embs = encoder(x).cpu() + embs.append(batch_embs.numpy()) + train_embs = np.concatenate(train_embs) + test_embs = np.concatenate(test_embs) + + os.makedirs(output_path, exist_ok=True) + np.save(os.path.join(output_path, "train_embeds.npy"), train_embs) + np.save(os.path.join(output_path, "test_embeds.npy"), test_embs) + + +if __name__ == "__main__": + args = parse_args() + main( + dataset_path=args.dataset_path, + output_path=args.output_path, + ckpt_path=args.ckpt, + batch_size=args.batch_size, + ) diff --git a/scripts/pretrain/_cache_heedb.py b/scripts/pretrain/_cache_heedb.py index 3d8d086..1f69255 100644 --- a/scripts/pretrain/_cache_heedb.py +++ b/scripts/pretrain/_cache_heedb.py @@ -11,4 +11,5 @@ dataset_path=DATASET_PATH, sampling_rate=100, split=split, + heedb_split_type="by-label", ) diff --git a/scripts/pretrain/run-heedb-normalizations.sh b/scripts/pretrain/run-heedb-normalizations.sh index f5b104c..e36a45a 100755 --- a/scripts/pretrain/run-heedb-normalizations.sh +++ b/scripts/pretrain/run-heedb-normalizations.sh @@ -3,6 +3,6 @@ set -e HEEDB_DATA=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/ecg-prototype-fm/data-preprocessing/heedb_normalizations +RUN_DIR=/opt/gpu_working/steven/ProtoSSL/data-preprocessing/heedb_normalizations python _compute_heedb_normalizations.py --dataset_path $HEEDB_DATA --output_path $RUN_DIR diff --git a/scripts/pretrain/run-pclr-heedb-pretrain.sh b/scripts/pretrain/run-pclr-heedb-pretrain.sh new file mode 100755 index 0000000..af8a5f6 --- /dev/null +++ b/scripts/pretrain/run-pclr-heedb-pretrain.sh @@ -0,0 +1,44 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=500gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl01 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/pclr-heedb-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +# set these env vars prior to executing this script +: "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using PRETRAIN_DATASET=$PRETRAIN_DATASET" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# experiment parameters +EXP_NAME="pclr-heedb" +BACKBONE=resnet18 +CONV=2D + +# pretrain via self supervised contrastive learning +python -m protossl.trainer \ + --pipeline-stage train-contraster \ + --config $REPO_ROOT/configs/pretrain-unsupervised-no-proto.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.max_epochs 100 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET diff --git a/scripts/pretrain/run-protossl-heedb-150-1050-pretrain.sh b/scripts/pretrain/run-protossl-heedb-150-1050-pretrain.sh new file mode 100755 index 0000000..850ab43 --- /dev/null +++ b/scripts/pretrain/run-protossl-heedb-150-1050-pretrain.sh @@ -0,0 +1,75 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=500gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl02 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protossl-heedb-150-1050-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +# set these env vars prior to executing this script +: "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using PRETRAIN_DATASET=$PRETRAIN_DATASET" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# toggle HEEDB label set +export USE_HEEDB_150=1 + +# experiment parameters +EXP_NAME="protossl-heedb-150-1050" +BACKBONE=resnet18 +CONV=2D + +# pretrain via self supervised prototype learning +python -m protossl.trainer \ + --pipeline-stage learn-prototypes \ + --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150-1050.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.max_epochs 100 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET + +# project in the pretraining dataset +python -m protossl.trainer \ + --pipeline-stage project-prototypes \ + --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150-1050.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes/latest/best.ckpt + +# python -m protossl.trainer \ +# --pipeline-stage train-classifier \ +# --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150-1050.yaml \ +# --model.backbone_type $BACKBONE \ +# --model.conv_type $CONV \ +# --trainer.logger.save_dir $RUN_DIR \ +# --trainer.logger.name $EXP_NAME \ +# --data.dataset_path $PRETRAIN_DATASET \ +# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes/latest/proj.ckpt + +# python _eval_probs.py \ +# --dataset-path $PRETRAIN_DATASET \ +# --probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ +# --output-path $RUN_DIR/$EXP_NAME + +# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy diff --git a/scripts/pretrain/run-protossl-heedb-150-pretrain.sh b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh new file mode 100755 index 0000000..0f99d2a --- /dev/null +++ b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh @@ -0,0 +1,75 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=500gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl02 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protossl-heedb-150-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +# set these env vars prior to executing this script +: "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using PRETRAIN_DATASET=$PRETRAIN_DATASET" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# toggle HEEDB label set +export USE_HEEDB_150=1 + +# experiment parameters +EXP_NAME="protossl-heedb-150" +BACKBONE=resnet18 +CONV=2D + +# pretrain via self supervised prototype learning +python -m protossl.trainer \ + --pipeline-stage learn-prototypes \ + --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.max_epochs 100 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET + +# project in the pretraining dataset +python -m protossl.trainer \ + --pipeline-stage project-prototypes \ + --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes/latest/best.ckpt + +python -m protossl.trainer \ + --pipeline-stage train-classifier \ + --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes/latest/proj.ckpt + +python _eval_probs.py \ +--dataset-path $PRETRAIN_DATASET \ +--probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy diff --git a/scripts/pretrain/run-protossl-heedb-pretrain-no-koleo.sh b/scripts/pretrain/run-protossl-heedb-pretrain-no-koleo.sh new file mode 100755 index 0000000..897bf99 --- /dev/null +++ b/scripts/pretrain/run-protossl-heedb-pretrain-no-koleo.sh @@ -0,0 +1,73 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=500gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl02 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/protossl-heedb-no-koleo-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +# set these env vars prior to executing this script +: "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using PRETRAIN_DATASET=$PRETRAIN_DATASET" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# experiment parameters +EXP_NAME="protossl-heedb-no-koleo" +BACKBONE=resnet18 +CONV=2D + +# pretrain via self supervised prototype learning +python -m protossl.trainer \ + --pipeline-stage learn-prototypes \ + --config $REPO_ROOT/configs/pretrain-unsupervised.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --model.model_kwargs '{"koleo_loss_weight": 0.0}' \ + --trainer.max_epochs 100 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET + +# project in the pretraining dataset +python -m protossl.trainer \ + --pipeline-stage project-prototypes \ + --config $REPO_ROOT/configs/pretrain-unsupervised.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes/latest/best.ckpt + +# python -m protossl.trainer \ +# --pipeline-stage train-classifier \ +# --config $REPO_ROOT/configs/pretrain-unsupervised.yaml \ +# --model.backbone_type $BACKBONE \ +# --model.conv_type $CONV \ +# --trainer.logger.save_dir $RUN_DIR \ +# --trainer.logger.name $EXP_NAME \ +# --data.dataset_path $PRETRAIN_DATASET \ +# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes/latest/proj.ckpt + +# python _eval_probs.py \ +# --dataset-path $PRETRAIN_DATASET \ +# --probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ +# --output-path $RUN_DIR/$EXP_NAME + +# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy diff --git a/scripts/pretrain/run-pass-heedb-pretrain.sh b/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh similarity index 92% rename from scripts/pretrain/run-pass-heedb-pretrain.sh rename to scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh index b87c07d..ae09c73 100755 --- a/scripts/pretrain/run-pass-heedb-pretrain.sh +++ b/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh @@ -3,8 +3,8 @@ set -e PRETRAIN_DATASET=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/ecg-prototype-fm/outputs -REPO_ROOT=/opt/gpudata/steven/ecg-prototype-fm +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL # set these env vars prior to executing this script : "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" @@ -19,7 +19,7 @@ cd $REPO_ROOT/scripts export HIGH_MEMORY=1 # experiment parameters -EXP_NAME="pass-pretrain-heedb" +EXP_NAME="protossl-heedb-w-attn" BACKBONE=resnet18 CONV=2D @@ -35,7 +35,7 @@ python -m protossl.trainer \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $PRETRAIN_DATASET -project in the pretraining dataset +# project in the pretraining dataset python -m protossl.trainer \ --pipeline-stage project-prototypes \ --config $REPO_ROOT/configs/pretrain-unsupervised.yaml \ diff --git a/scripts/pretrain/run-pass-heedb-pretrain-no-attn.sh b/scripts/pretrain/run-protossl-heedb-pretrain.sh similarity index 93% rename from scripts/pretrain/run-pass-heedb-pretrain-no-attn.sh rename to scripts/pretrain/run-protossl-heedb-pretrain.sh index 2255794..a05532f 100755 --- a/scripts/pretrain/run-pass-heedb-pretrain-no-attn.sh +++ b/scripts/pretrain/run-protossl-heedb-pretrain.sh @@ -3,8 +3,8 @@ set -e PRETRAIN_DATASET=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/ecg-prototype-fm/outputs -REPO_ROOT=/opt/gpudata/steven/ecg-prototype-fm +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL # set these env vars prior to executing this script : "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" @@ -19,7 +19,7 @@ cd $REPO_ROOT/scripts export HIGH_MEMORY=1 # experiment parameters -EXP_NAME="pass-pretrain-heedb-no-attn" +EXP_NAME="protossl-heedb" BACKBONE=resnet18 CONV=2D diff --git a/scripts/pretrain/run-supproto-heedb-150-1050-pretrain.sh b/scripts/pretrain/run-supproto-heedb-150-1050-pretrain.sh new file mode 100755 index 0000000..75b9562 --- /dev/null +++ b/scripts/pretrain/run-supproto-heedb-150-1050-pretrain.sh @@ -0,0 +1,81 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=500gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl02 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/supproto-heedb-150-1050-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +# set these env vars prior to executing this script +: "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using PRETRAIN_DATASET=$PRETRAIN_DATASET" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# toggle HEEDB label set +export USE_HEEDB_150=1 + +# experiment parameters +EXP_NAME="supproto-heedb-150-1050" +BACKBONE=resnet18 +CONV=2D + +# pretrain via label supervised prototype learning +python -m protossl.trainer \ + --pipeline-stage learn-prototypes-supervised \ + --config $REPO_ROOT/configs/pretrain-supervised-heedb-150-1050.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.max_epochs 100 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +# project in the pretraining dataset +python -m protossl.trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/pretrain-supervised-heedb-150-1050.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +# python -m protossl.trainer \ +# --pipeline-stage train-classifier \ +# --config $REPO_ROOT/configs/pretrain-supervised-heedb-150-1050.yaml \ +# --model.backbone_type $BACKBONE \ +# --model.conv_type $CONV \ +# --trainer.logger.save_dir $RUN_DIR \ +# --trainer.logger.name $EXP_NAME \ +# --data.dataset_path $PRETRAIN_DATASET \ +# --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt \ +# --data.num_workers 8 \ +# --data.prefetch_factor 4 + +# python _eval_probs.py \ +# --dataset-path $PRETRAIN_DATASET \ +# --probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ +# --output-path $RUN_DIR/$EXP_NAME + +# cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy diff --git a/scripts/pretrain/run-supproto-heedb-150-pretrain.sh b/scripts/pretrain/run-supproto-heedb-150-pretrain.sh new file mode 100755 index 0000000..f71c313 --- /dev/null +++ b/scripts/pretrain/run-supproto-heedb-150-pretrain.sh @@ -0,0 +1,81 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=500gb +#SBATCH --gpus-per-node=1 +#SBATCH --nodes=1 +#SBATCH -w kg35-nvl02 +#SBATCH --ntasks-per-node=1 +#SBATCH --time=0 +#SBATCH --output /home/songs1/slurm-logs/supproto-heedb-150-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL + +# set these env vars prior to executing this script +: "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" +: "${RUN_DIR:?Env var RUN_DIR must be set prior to script execution}" +: "${REPO_ROOT:?Env var REPO_ROOT must be set prior to script execution}" +echo "Using PRETRAIN_DATASET=$PRETRAIN_DATASET" +echo "Using RUN_DIR=$RUN_DIR" +echo "Using REPO_ROOT=$REPO_ROOT" +cd $REPO_ROOT/scripts + +# submit job with 500GB of memory +export HIGH_MEMORY=1 + +# toggle HEEDB label set +export USE_HEEDB_150=1 + +# experiment parameters +EXP_NAME="supproto-heedb-150" +BACKBONE=resnet18 +CONV=2D + +# pretrain via label supervised prototype learning +python -m protossl.trainer \ + --pipeline-stage learn-prototypes-supervised \ + --config $REPO_ROOT/configs/pretrain-supervised-heedb-150.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.max_epochs 100 \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +# project in the pretraining dataset +python -m protossl.trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/pretrain-supervised-heedb-150.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/latest/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python -m protossl.trainer \ + --pipeline-stage train-classifier \ + --config $REPO_ROOT/configs/pretrain-supervised-heedb-150.yaml \ + --model.backbone_type $BACKBONE \ + --model.conv_type $CONV \ + --trainer.logger.save_dir $RUN_DIR \ + --trainer.logger.name $EXP_NAME \ + --data.dataset_path $PRETRAIN_DATASET \ + --model.pretrained_weights $RUN_DIR/$EXP_NAME/project-prototypes-supervised/latest/proj.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 + +python _eval_probs.py \ +--dataset-path $PRETRAIN_DATASET \ +--probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +cp $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy diff --git a/scripts/pretrain/run-prosup-heedb-pretrain.sh b/scripts/pretrain/run-supproto-heedb-pretrain.sh similarity index 93% rename from scripts/pretrain/run-prosup-heedb-pretrain.sh rename to scripts/pretrain/run-supproto-heedb-pretrain.sh index e438b74..931a24e 100755 --- a/scripts/pretrain/run-prosup-heedb-pretrain.sh +++ b/scripts/pretrain/run-supproto-heedb-pretrain.sh @@ -3,8 +3,8 @@ set -e PRETRAIN_DATASET=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/ecg-prototype-fm/outputs -REPO_ROOT=/opt/gpudata/steven/ecg-prototype-fm +RUN_DIR=/opt/gpu_working/steven/protossl-ecg-outputs +REPO_ROOT=/opt/gpu_working/steven/ProtoSSL # set these env vars prior to executing this script : "${PRETRAIN_DATASET:?Env var PRETRAIN_DATASET must be set prior to script execution}" @@ -15,10 +15,11 @@ echo "Using RUN_DIR=$RUN_DIR" echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts +# submit job with 500GB of memory export HIGH_MEMORY=1 # experiment parameters -EXP_NAME="prosup-pretrain-heedb" +EXP_NAME="supproto-heedb" BACKBONE=resnet18 CONV=2D diff --git a/scripts/queue-experiments.sh b/scripts/queue-experiments.sh index a08c60f..5d60e5c 100755 --- a/scripts/queue-experiments.sh +++ b/scripts/queue-experiments.sh @@ -28,7 +28,7 @@ export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL for seed in "${SEEDS[@]}"; do echo "SEED: $seed" export SEED=$seed - export BASE_OUTPUT_DIR=/opt/gpu_working/steven/new/protossl-outputs-seed"$SEED" + export BASE_OUTPUT_DIR=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal/protossl-outputs-seed"$SEED" for dataset in "${DATASETS[@]}"; do echo "DATASET: $dataset" @@ -46,25 +46,37 @@ for seed in "${SEEDS[@]}"; do echo $cache_id submit_job "$suffix" 1-run-blackbox-direct.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 2-run-labsup-proto-direct.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 2-run-supproto-direct.sh "--dependency=afterok:$cache_id" submit_job "$suffix" 3-run-protossl-heedb-pila.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 4-run-labsup-proto-heedb-rila.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 4-run-supproto-heedb-rila.sh "--dependency=afterok:$cache_id" submit_job "$suffix" 5-1-run-ecgfounder-logreg.sh # does not depend on same 100 Hz cache (takes 500 Hz) submit_job "$suffix" 5-2-run-stmem-logreg.sh # does not depend on same 100 Hz cache (takes 250 Hz) + submit_job "$suffix" 5-3-run-pclr-logreg.sh "--dependency=afterok:$cache_id" + + if [ "$dataset" == "zzu" ] && [ "$suffix" == "" ]; then + echo "ZZU full scale, larger prototype bank" + submit_job "$suffix" 3-run-protossl-heedb-150-pila.sh "--dependency=afterok:$cache_id" + fi if [ "$dataset" == "echonext" ]; then echo "Ablations" # prototypes per label - submit_job "$suffix" 2-z-run-labsup-proto-direct-7ppl.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 2-z-run-labsup-proto-direct-28ppl.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 2-z-run-supproto-direct-7ppl.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 2-z-run-supproto-direct-28ppl.sh "--dependency=afterok:$cache_id" submit_job "$suffix" 3-z-run-protossl-heedb-pila-7ppl.sh "--dependency=afterok:$cache_id" submit_job "$suffix" 3-z-run-protossl-heedb-pila-28ppl.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 4-z-run-labsup-proto-heedb-rila-7ppl.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 4-z-run-labsup-proto-heedb-rila-28ppl.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 4-z-run-supproto-heedb-rila-7ppl.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 4-z-run-supproto-heedb-rila-28ppl.sh "--dependency=afterok:$cache_id" + + # pretraining over HEEDB-150 labels, 7500 (150 labels x 50 ppl) and 1050 prototype banks + submit_job "$suffix" 3-run-protossl-heedb-150-pila.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 4-run-supproto-heedb-150-rila.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 3-run-protossl-heedb-150-1050-pila.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 4-run-supproto-heedb-150-1050-rila.sh "--dependency=afterok:$cache_id" # supproto no-proj - submit_job "$suffix" 4-y-run-labsup-proto-heedb-pila.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 4-y-run-supproto-heedb-pila.sh "--dependency=afterok:$cache_id" # assignment method submit_job "$suffix" 6-run-protossl-heedb-pia.sh "--dependency=afterok:$cache_id" @@ -72,8 +84,9 @@ for seed in "${SEEDS[@]}"; do # no assignment submit_job "$suffix" 3-z-run-protossl-heedb-pila-83ppl.sh "--dependency=afterok:$cache_id" submit_job "$suffix" 6-z-run-protossl-heedb-pia-83ppl.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 7-run-protossl-heedb-pit.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 8-run-protossl-heedb-pip.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 7-1-run-protossl-heedb-pit.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 7-2-run-protossl-heedb-pip.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 7-3-run-protossl-heedb-latent.sh "--dependency=afterok:$cache_id" # start with pretrained encoder patches_id=$(submit_job "$suffix" 9-0-run-ecgfounder-patches.sh) diff --git a/scripts/queue-pclr-logreg.sh b/scripts/queue-pclr-logreg.sh new file mode 100644 index 0000000..11355b8 --- /dev/null +++ b/scripts/queue-pclr-logreg.sh @@ -0,0 +1,52 @@ +#!/bin/bash + +# PCLR (HEEDB-pretrained, no prototypes) linear probe + +set -e + +# declare -a DATASETS=("echonext" "ptbxl" "cinc" "mimic" "zzu" "code15") +declare -a DATASETS=("ptbxl" "cinc" "mimic" "zzu" "code15") +declare -a SEEDS=(42 67 70 73 99) +declare -A DATASET_DIRS=( + ["echonext"]="echonext" + ["ptbxl"]="ptb-xl" + ["cinc"]="cinc-2020" + ["mimic"]="mimic-iv-ecg" + ["zzu"]="zzu-pecg" + ["code15"]="code15" +) +declare -a SUFFIXES_echonext=("" "-32k" "-16k" "-8k" "-4k" "-2k" "-1k" "-512" "-256") +declare -a SUFFIXES_ptbxl=("" "-8k" "-4k" "-2k" "-1k" "-512" "-256") +declare -a SUFFIXES_code15=("" "-32k" "-16k" "-8k" "-4k" "-2k" "-1k" "-512" "-256") +declare -a SUFFIXES_cinc=("" "-4k" "-2k" "-1k" "-512" "-256") +declare -a SUFFIXES_mimic=("" "-32k" "-16k" "-8k" "-4k" "-2k" "-1k" "-512" "-256") +declare -a SUFFIXES_zzu=("" "-4k" "-2k" "-1k" "-512" "-256") + +source _submit_job.sh + +export BASE_ECG_PATH=/opt/gpudata/ecg +export REPO_ROOT=/opt/gpu_working/steven/ProtoSSL +# 5-3-run-pclr-logreg.sh finds the pretrained model at $RUN_DIR/../../pclr-heedb +BASE_OUTPUT_ROOT=/opt/gpu_working/steven/protossl-ecg-outputs-rebuttal + +for seed in "${SEEDS[@]}"; do + echo "SEED: $seed" + export SEED=$seed + export BASE_OUTPUT_DIR="$BASE_OUTPUT_ROOT/protossl-outputs-seed$SEED" + + for dataset in "${DATASETS[@]}"; do + echo "DATASET: $dataset" + ds_dir="${DATASET_DIRS[$dataset]}" + export BASE_RUN_DIR="$BASE_OUTPUT_DIR/runs-$dataset" + export BASE_DATASET_PATH="$BASE_ECG_PATH/$ds_dir" + + array_name="SUFFIXES_${dataset}[@]" + suffixes=("${!array_name}") + + for suffix in "${suffixes[@]}"; do + echo "SUFFIX: $suffix" + # 100 Hz caches for all these datasets/sizes already exist from the main experiments + submit_job "$suffix" 5-3-run-pclr-logreg.sh + done + done +done diff --git a/user-study/images b/user-study/images index e18cdd0..c838491 160000 --- a/user-study/images +++ b/user-study/images @@ -1 +1 @@ -Subproject commit e18cdd08939aad7ad673938e818c97b75494ec65 +Subproject commit c8384916f70c3de33aceaeefa0810251df81e1a3 diff --git a/user-study/prepare_samples.ipynb b/user-study/prepare_samples.ipynb index 74db6a8..4740f9b 100644 --- a/user-study/prepare_samples.ipynb +++ b/user-study/prepare_samples.ipynb @@ -20,6 +20,17 @@ { "cell_type": "code", "execution_count": 1, + "id": "1d5cf381", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"CACHE_DIR\"] = \"/home/songs1/.cache/protossl_cache\" # for working in dev container" + ] + }, + { + "cell_type": "code", + "execution_count": 2, "id": "ee62869b", "metadata": {}, "outputs": [], @@ -43,8 +54,8 @@ "ptbxl_path = \"/opt/gpudata/ecg/ptb-xl\"\n", "labels = [\"AMI\", \"CRBBB\", \"CLBBB\", \"PVC\"]\n", "cases_per_label = 5 # CPL\n", - "protossl = Path(\"/opt/gpu_working/steven/protossl-outputs/runs-ptbxl/protossl-heedb-pila\")\n", - "protoecgnet = Path(\"/opt/gpu_working/steven/protossl-outputs/runs-ptbxl/labsup-proto-direct\")\n", + "protossl = Path(\"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed42/runs-ptbxl/protossl-heedb-pila\")\n", + "protoecgnet = Path(\"/opt/gpu_working/steven/protossl-ecg-outputs/experiments-seed42/runs-ptbxl/supproto-direct\")\n", "prototypes_per_label = 14 # the PPL used in above models\n", "\n", "rng = np.random.default_rng(12)\n", @@ -54,7 +65,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 3, "id": "370fba08", "metadata": {}, "outputs": [ @@ -64,11 +75,11 @@ "text": [ "=================load_cached_data==================\n", "Dataset parameters: (/opt/gpudata/ecg/ptb-xl, test, 100Hz)\n", - "Cached dataset loaded from /home/songs1/.cache/pass_pclr_cache/9a71164f.pt\n", + "Cached dataset loaded from /home/songs1/.cache/protossl_cache/9a71164f.pt\n", "===================================================\n", "=================load_cached_data==================\n", "Dataset parameters: (/opt/gpudata/ecg/ptb-xl, train, 100Hz)\n", - "Cached dataset loaded from /home/songs1/.cache/pass_pclr_cache/52f12d62.pt\n", + "Cached dataset loaded from /home/songs1/.cache/protossl_cache/52f12d62.pt\n", "===================================================\n" ] } @@ -83,7 +94,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 4, "id": "cc8ea4d4", "metadata": {}, "outputs": [], @@ -94,7 +105,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 5, "id": "29aae7a7", "metadata": {}, "outputs": [ @@ -258,7 +269,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 6, "id": "c2369fa0", "metadata": {}, "outputs": [], @@ -299,7 +310,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 7, "id": "868bbd9f", "metadata": {}, "outputs": [], @@ -314,7 +325,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 8, "id": "30c1c262", "metadata": {}, "outputs": [], @@ -599,7 +610,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 9, "id": "cd9cf6c1", "metadata": {}, "outputs": [ @@ -607,7 +618,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|██████████| 20/20 [02:10<00:00, 6.54s/it]\n" + "100%|██████████| 20/20 [01:26<00:00, 4.32s/it]\n" ] } ], @@ -671,7 +682,8 @@ " right_ecg_window=b_meta[label_name][\"prototype_ecg_window\"][j],\n", " right_title=model_b_name,\n", " )\n", - " global_fig.savefig(f\"images/{study_idx:02d}_global_proto.png\", dpi=300)\n", + " # global_fig.savefig(f\"images/{study_idx:02d}_global_proto.png\", dpi=300)\n", + " global_fig.savefig(f\"images-pdf/{study_idx:02d}_global_proto.pdf\")\n", " a_fig = plot_compare(\n", " sup_title=f\"Case {study_idx}: {label_clean}\",\n", " # case on left\n", @@ -685,7 +697,8 @@ " # annotate similarity\n", " annotation=f\"{a_meta[label_name]['similarity'][j]:0.3f} Similarity\",\n", " )\n", - " a_fig.savefig(f\"images/{study_idx:02d}_case_proto_A.png\", dpi=300)\n", + " # a_fig.savefig(f\"images/{study_idx:02d}_case_proto_A.png\", dpi=300)\n", + " a_fig.savefig(f\"images-pdf/{study_idx:02d}_case_proto_A.pdf\")\n", " b_fig = plot_compare(\n", " sup_title=f\"Case {study_idx}: {label_clean}\",\n", " # case on left\n", @@ -699,7 +712,8 @@ " # annotate similarity\n", " annotation=f\"{b_meta[label_name]['similarity'][j]:0.3f} Similarity\",\n", " )\n", - " b_fig.savefig(f\"images/{study_idx:02d}_case_proto_B.png\", dpi=300)\n", + " # b_fig.savefig(f\"images/{study_idx:02d}_case_proto_B.png\", dpi=300)\n", + " b_fig.savefig(f\"images-pdf/{study_idx:02d}_case_proto_B.pdf\")\n", "\n", " plt.close(\"all\")\n", "\n", @@ -709,7 +723,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 10, "id": "f32a14ec", "metadata": {}, "outputs": [], @@ -724,7 +738,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 11, "id": "e5088edc", "metadata": {}, "outputs": [ @@ -1265,7 +1279,7 @@ "19 0.948079 " ] }, - "execution_count": 10, + "execution_count": 11, "metadata": {}, "output_type": "execute_result" } @@ -1276,7 +1290,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 12, "id": "ed4df38a", "metadata": {}, "outputs": [ @@ -1433,7 +1447,7 @@ "19 0.948079 " ] }, - "execution_count": 11, + "execution_count": 12, "metadata": {}, "output_type": "execute_result" } @@ -1445,7 +1459,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 13, "id": "1b0b251b", "metadata": {}, "outputs": [],