From 32e2790c097fb87496f42d646042198f65673130 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Tue, 21 Jul 2026 10:41:55 -0500 Subject: [PATCH 01/33] redo audio --- .gitignore | 1 + configs/audio/pretrain-supervised-v2.yaml | 60 +++++++ configs/audio/pretrain-unsupervised-v2.yaml | 63 +++++++ configs/audio/pretrain-unsupervised.yaml | 2 +- protossl/models/_prototype_contraster.py | 4 +- protossl/trainer.py | 12 +- scripts/audio/3-run-protossl-audioset-pila.sh | 160 +++++++++--------- .../audio/4-run-labsup-proto-audioset-rila.sh | 160 +++++++++--------- scripts/audio/queue-experiments.sh | 14 +- ...pass.sh => run-pretrain-audio-protossl.sh} | 6 +- ...osup.sh => run-pretrain-audio-supproto.sh} | 20 +-- .../voxceleb-2-run-labsup-proto-direct.sh | 64 +++---- .../voxceleb-3-run-protossl-audioset-pila.sh | 152 ++++++++--------- ...xceleb-4-run-labsup-proto-audioset-rila.sh | 152 ++++++++--------- user-study/images | 2 +- 15 files changed, 501 insertions(+), 371 deletions(-) create mode 100644 configs/audio/pretrain-supervised-v2.yaml create mode 100644 configs/audio/pretrain-unsupervised-v2.yaml rename scripts/audio/{run-pretrain-audio-pass.sh => run-pretrain-audio-protossl.sh} (86%) rename scripts/audio/{run-pretrain-audio-prosup.sh => run-pretrain-audio-supproto.sh} (73%) diff --git a/.gitignore b/.gitignore index 2c32da1..dcdba4c 100644 --- a/.gitignore +++ b/.gitignore @@ -219,3 +219,4 @@ OLD/ slurm-logs/ results/figs/ ecgfounder-checkpoint/ +configs/multi-arm-heedb/wandb diff --git a/configs/audio/pretrain-supervised-v2.yaml b/configs/audio/pretrain-supervised-v2.yaml new file mode 100644 index 0000000..52c55a0 --- /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: 5 + 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-unsupervised-v2.yaml b/configs/audio/pretrain-unsupervised-v2.yaml new file mode 100644 index 0000000..6ac7290 --- /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: 5 + 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..80589a1 100644 --- a/configs/audio/pretrain-unsupervised.yaml +++ b/configs/audio/pretrain-unsupervised.yaml @@ -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/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/trainer.py b/protossl/trainer.py index 02c9e2a..5d0e89e 100644 --- a/protossl/trainer.py +++ b/protossl/trainer.py @@ -54,6 +54,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() @@ -94,6 +95,7 @@ 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 + data_wrapper_kwargs: dict = self.hparams.data_wrapper_kwargs # type: ignore wrap_contrastive = pipeline_stage == "learn-prototypes" if stage == "fit": @@ -116,11 +118,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 +143,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( diff --git a/scripts/audio/3-run-protossl-audioset-pila.sh b/scripts/audio/3-run-protossl-audioset-pila.sh index eb0e8c8..dbe1ad1 100755 --- a/scripts/audio/3-run-protossl-audioset-pila.sh +++ b/scripts/audio/3-run-protossl-audioset-pila.sh @@ -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="$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 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 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 index a696997..21f8f70 100755 --- a/scripts/audio/4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/4-run-labsup-proto-audioset-rila.sh @@ -28,48 +28,48 @@ 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 +PRETRAIN_RUN="$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 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 \ @@ -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/queue-experiments.sh b/scripts/audio/queue-experiments.sh index ff3acb4..5546323 100644 --- a/scripts/audio/queue-experiments.sh +++ b/scripts/audio/queue-experiments.sh @@ -1,14 +1,14 @@ #!/bin/bash scripts=( - "1-run-blackbox-direct.sh" - "2-run-labsup-proto-direct.sh" + # "1-run-blackbox-direct.sh" + # "2-run-labsup-proto-direct.sh" "3-run-protossl-audioset-pila.sh" "4-run-labsup-proto-audioset-rila.sh" ) voxceleb_scripts=( - "1-run-blackbox-direct.sh" + # "1-run-blackbox-direct.sh" "voxceleb-2-run-labsup-proto-direct.sh" "voxceleb-3-run-protossl-audioset-pila.sh" "voxceleb-4-run-labsup-proto-audioset-rila.sh" @@ -18,7 +18,7 @@ voxceleb_scripts=( 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 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 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 +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 86% rename from scripts/audio/run-pretrain-audio-pass.sh rename to scripts/audio/run-pretrain-audio-protossl.sh index 824001c..e67b924 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,7 @@ 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 diff --git a/scripts/audio/run-pretrain-audio-prosup.sh b/scripts/audio/run-pretrain-audio-supproto.sh similarity index 73% rename from scripts/audio/run-pretrain-audio-prosup.sh rename to scripts/audio/run-pretrain-audio-supproto.sh index 96c62ac..2cd2367 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,19 @@ 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}' echo "======================================" echo "Stage 2: project prototypes supervised" @@ -46,7 +46,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 index fc716a0..719d1fe 100755 --- a/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh +++ b/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh @@ -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}" @@ -33,40 +33,40 @@ 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 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 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 +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 +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.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-3-run-protossl-audioset-pila.sh b/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh index 7158ae9..def1405 100755 --- a/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh +++ b/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh @@ -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="$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 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 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 index 5fd8ab9..89e2e7a 100755 --- a/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh @@ -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="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 +PRETRAIN_RUN="$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 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 \ @@ -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/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 From 365a92028cf30a294ee34b376652611574027b5a Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 23 Jul 2026 11:41:13 -0500 Subject: [PATCH 02/33] correct pretraining max_epochs --- configs/audio/pretrain-supervised-v2.yaml | 2 +- configs/audio/pretrain-supervised.yaml | 2 +- configs/audio/pretrain-unsupervised-v2.yaml | 2 +- configs/audio/pretrain-unsupervised.yaml | 2 +- scripts/audio/3-run-protossl-audioset-pila.sh | 4 ++-- scripts/audio/4-run-labsup-proto-audioset-rila.sh | 4 ++-- scripts/audio/queue-experiments.sh | 8 ++++---- scripts/audio/run-pretrain-audio-protossl.sh | 3 ++- scripts/audio/run-pretrain-audio-supproto.sh | 3 ++- scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh | 4 ++-- .../audio/voxceleb-4-run-labsup-proto-audioset-rila.sh | 4 ++-- 11 files changed, 20 insertions(+), 18 deletions(-) diff --git a/configs/audio/pretrain-supervised-v2.yaml b/configs/audio/pretrain-supervised-v2.yaml index 52c55a0..0b91ab6 100644 --- a/configs/audio/pretrain-supervised-v2.yaml +++ b/configs/audio/pretrain-supervised-v2.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-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 index 6ac7290..a91c8f1 100644 --- a/configs/audio/pretrain-unsupervised-v2.yaml +++ b/configs/audio/pretrain-unsupervised-v2.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.yaml b/configs/audio/pretrain-unsupervised.yaml index 80589a1..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 diff --git a/scripts/audio/3-run-protossl-audioset-pila.sh b/scripts/audio/3-run-protossl-audioset-pila.sh index dbe1ad1..c6f72b8 100755 --- a/scripts/audio/3-run-protossl-audioset-pila.sh +++ b/scripts/audio/3-run-protossl-audioset-pila.sh @@ -28,14 +28,14 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="protossl-audioset-pila" -PRETRAIN_RUN="$RUN_DIR/../../protossl-audioset-v2" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2.BAK") 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 \ + --model.n_prototypes 6324 \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.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 index 21f8f70..4ab7b29 100755 --- a/scripts/audio/4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/4-run-labsup-proto-audioset-rila.sh @@ -28,14 +28,14 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="labsup-proto-audioset-rila" -PRETRAIN_RUN="$RUN_DIR/../../supproto-audioset-v2" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2.BAK") 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 \ + --model.n_prototypes 6324 \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ diff --git a/scripts/audio/queue-experiments.sh b/scripts/audio/queue-experiments.sh index 5546323..f401a72 100644 --- a/scripts/audio/queue-experiments.sh +++ b/scripts/audio/queue-experiments.sh @@ -9,7 +9,7 @@ scripts=( voxceleb_scripts=( # "1-run-blackbox-direct.sh" - "voxceleb-2-run-labsup-proto-direct.sh" + # "voxceleb-2-run-labsup-proto-direct.sh" "voxceleb-3-run-protossl-audioset-pila.sh" "voxceleb-4-run-labsup-proto-audioset-rila.sh" ) @@ -18,7 +18,7 @@ voxceleb_scripts=( 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-outputs/experiments/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-outputs/experiments/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,7 +36,7 @@ done # voxceleb1 (these take even longer to run) export DATASET_PATH=/opt/gpudata/audio/VoxCeleb1 -export RUN_DIR=/opt/gpu_working/steven/protossl-audio-outputs/experiments/runs-voxceleb +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 diff --git a/scripts/audio/run-pretrain-audio-protossl.sh b/scripts/audio/run-pretrain-audio-protossl.sh index e67b924..f88908b 100644 --- a/scripts/audio/run-pretrain-audio-protossl.sh +++ b/scripts/audio/run-pretrain-audio-protossl.sh @@ -37,4 +37,5 @@ srun python -m protossl.trainer \ --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-supproto.sh b/scripts/audio/run-pretrain-audio-supproto.sh index 2cd2367..fde1e22 100644 --- a/scripts/audio/run-pretrain-audio-supproto.sh +++ b/scripts/audio/run-pretrain-audio-supproto.sh @@ -38,7 +38,8 @@ srun python -m protossl.trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --model.model_kwargs '{"use_default_weights": True}' + --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" diff --git a/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh b/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh index def1405..de0e0f8 100755 --- a/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh +++ b/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh @@ -32,14 +32,14 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="protossl-audioset-pila" -PRETRAIN_RUN="$RUN_DIR/../../protossl-audioset-v2" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2.BAK") 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 \ + --model.n_prototypes 6324 \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.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 index 89e2e7a..036d455 100755 --- a/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh @@ -32,14 +32,14 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="labsup-proto-audioset-rila" -PRETRAIN_RUN="$RUN_DIR/../../supproto-audioset-v2" +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2.BAK") 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 \ + --model.n_prototypes 6324 \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ From d3de0a598d51ce2435f72a9decdb9fc2c6f45423 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 23 Jul 2026 14:11:23 -0500 Subject: [PATCH 03/33] toggle for heedb 150 labels --- protossl/defines.py | 156 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) diff --git a/protossl/defines.py b/protossl/defines.py index 7a8a2a3..ae6a6de 100644 --- a/protossl/defines.py +++ b/protossl/defines.py @@ -523,6 +523,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", From 42ff8783440944897c9e7eab3012fa419c8670fe Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 23 Jul 2026 14:35:55 -0500 Subject: [PATCH 04/33] add new run files, minor cleanup of file names --- .github/CODEOWNERS | 16 +--- configs/pretrain-supervised-heedb-150.yaml | 55 +++++++++++++ configs/pretrain-unsupervised-heedb-150.yaml | 56 +++++++++++++ .../run-protossl-heedb-150-pretrain.sh | 76 +++++++++++++++++ ...=> run-protossl-heedb-pretrain-no-attn.sh} | 0 ...rain.sh => run-protossl-heedb-pretrain.sh} | 0 .../run-supproto-heedb-150-pretrain.sh | 81 +++++++++++++++++++ ...rain.sh => run-supproto-heedb-pretrain.sh} | 1 + 8 files changed, 271 insertions(+), 14 deletions(-) create mode 100644 configs/pretrain-supervised-heedb-150.yaml create mode 100644 configs/pretrain-unsupervised-heedb-150.yaml create mode 100755 scripts/pretrain/run-protossl-heedb-150-pretrain.sh rename scripts/pretrain/{run-pass-heedb-pretrain-no-attn.sh => run-protossl-heedb-pretrain-no-attn.sh} (100%) rename scripts/pretrain/{run-pass-heedb-pretrain.sh => run-protossl-heedb-pretrain.sh} (100%) create mode 100755 scripts/pretrain/run-supproto-heedb-150-pretrain.sh rename scripts/pretrain/{run-prosup-heedb-pretrain.sh => run-supproto-heedb-pretrain.sh} (98%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index fdf0780..eac1b2b 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 @@ -95,12 +88,7 @@ scripts/ecg-fms/_compute_stmem_embeddings.py @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/README.md @StevenSong @sahilsethi0105 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.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/scripts/pretrain/run-protossl-heedb-150-pretrain.sh b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh new file mode 100755 index 0000000..d7f920e --- /dev/null +++ b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh @@ -0,0 +1,76 @@ +#!/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=/home/songs1/protossl-ecg-outputs-rebuttal +REPO_ROOT=/home/songs1/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 \ + --model.model_kwargs '{"do_softmax": True, "do_weighted_sum": True}' \ + --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-pass-heedb-pretrain-no-attn.sh b/scripts/pretrain/run-protossl-heedb-pretrain-no-attn.sh similarity index 100% rename from scripts/pretrain/run-pass-heedb-pretrain-no-attn.sh rename to scripts/pretrain/run-protossl-heedb-pretrain-no-attn.sh diff --git a/scripts/pretrain/run-pass-heedb-pretrain.sh b/scripts/pretrain/run-protossl-heedb-pretrain.sh similarity index 100% rename from scripts/pretrain/run-pass-heedb-pretrain.sh rename to scripts/pretrain/run-protossl-heedb-pretrain.sh 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..1e32a96 --- /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/protossl-heedb-150-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/home/songs1/protossl-ecg-outputs-rebuttal +REPO_ROOT=/home/songs1/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 98% rename from scripts/pretrain/run-prosup-heedb-pretrain.sh rename to scripts/pretrain/run-supproto-heedb-pretrain.sh index e438b74..44833e9 100755 --- a/scripts/pretrain/run-prosup-heedb-pretrain.sh +++ b/scripts/pretrain/run-supproto-heedb-pretrain.sh @@ -15,6 +15,7 @@ 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 From 424daf93423f588e876a39ec108137e3204a5ff8 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 23 Jul 2026 22:22:36 -0500 Subject: [PATCH 05/33] typo in heedb 150 things --- protossl/datasets/streaming_loaders/_streaming_base.py | 3 +++ protossl/defines.py | 2 +- scripts/pretrain/run-supproto-heedb-150-pretrain.sh | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) 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 ae6a6de..37b77c6 100644 --- a/protossl/defines.py +++ b/protossl/defines.py @@ -80,7 +80,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 diff --git a/scripts/pretrain/run-supproto-heedb-150-pretrain.sh b/scripts/pretrain/run-supproto-heedb-150-pretrain.sh index 1e32a96..762717c 100755 --- a/scripts/pretrain/run-supproto-heedb-150-pretrain.sh +++ b/scripts/pretrain/run-supproto-heedb-150-pretrain.sh @@ -7,7 +7,7 @@ #SBATCH -w kg35-nvl02 #SBATCH --ntasks-per-node=1 #SBATCH --time=0 -#SBATCH --output /home/songs1/slurm-logs/protossl-heedb-150-%j.out +#SBATCH --output /home/songs1/slurm-logs/supproto-heedb-150-%j.out set -e From e09fe49b60433089b237a530a9456443b6d7ea61 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 23 Jul 2026 22:55:49 -0500 Subject: [PATCH 06/33] typo, no attention for default protossl --- scripts/pretrain/run-protossl-heedb-150-pretrain.sh | 3 +-- ...etrain-no-attn.sh => run-protossl-heedb-pretrain-w-attn.sh} | 3 ++- scripts/pretrain/run-protossl-heedb-pretrain.sh | 3 +-- 3 files changed, 4 insertions(+), 5 deletions(-) rename scripts/pretrain/{run-protossl-heedb-pretrain-no-attn.sh => run-protossl-heedb-pretrain-w-attn.sh} (95%) diff --git a/scripts/pretrain/run-protossl-heedb-150-pretrain.sh b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh index d7f920e..09e8db3 100755 --- a/scripts/pretrain/run-protossl-heedb-150-pretrain.sh +++ b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh @@ -41,13 +41,12 @@ python -m protossl.trainer \ --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150.yaml \ --model.backbone_type $BACKBONE \ --model.conv_type $CONV \ - --model.model_kwargs '{"do_softmax": True, "do_weighted_sum": True}' \ --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 +# project in the pretraining dataset python -m protossl.trainer \ --pipeline-stage project-prototypes \ --config $REPO_ROOT/configs/pretrain-unsupervised-heedb-150.yaml \ diff --git a/scripts/pretrain/run-protossl-heedb-pretrain-no-attn.sh b/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh similarity index 95% rename from scripts/pretrain/run-protossl-heedb-pretrain-no-attn.sh rename to scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh index 2255794..fda90f2 100755 --- a/scripts/pretrain/run-protossl-heedb-pretrain-no-attn.sh +++ b/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh @@ -19,7 +19,7 @@ cd $REPO_ROOT/scripts export HIGH_MEMORY=1 # experiment parameters -EXP_NAME="pass-pretrain-heedb-no-attn" +EXP_NAME="pass-pretrain-heedb-w-attn" BACKBONE=resnet18 CONV=2D @@ -29,6 +29,7 @@ python -m protossl.trainer \ --config $REPO_ROOT/configs/pretrain-unsupervised.yaml \ --model.backbone_type $BACKBONE \ --model.conv_type $CONV \ + --model.model_kwargs '{"do_softmax": True, "do_weighted_sum": True}' \ --trainer.max_epochs 100 \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ diff --git a/scripts/pretrain/run-protossl-heedb-pretrain.sh b/scripts/pretrain/run-protossl-heedb-pretrain.sh index b87c07d..2255794 100755 --- a/scripts/pretrain/run-protossl-heedb-pretrain.sh +++ b/scripts/pretrain/run-protossl-heedb-pretrain.sh @@ -19,7 +19,7 @@ cd $REPO_ROOT/scripts export HIGH_MEMORY=1 # experiment parameters -EXP_NAME="pass-pretrain-heedb" +EXP_NAME="pass-pretrain-heedb-no-attn" BACKBONE=resnet18 CONV=2D @@ -29,7 +29,6 @@ python -m protossl.trainer \ --config $REPO_ROOT/configs/pretrain-unsupervised.yaml \ --model.backbone_type $BACKBONE \ --model.conv_type $CONV \ - --model.model_kwargs '{"do_softmax": True, "do_weighted_sum": True}' \ --trainer.max_epochs 100 \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ From 0eaeab038d8e8dc0a87402a3e0146f379e75bbf8 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 23 Jul 2026 22:57:37 -0500 Subject: [PATCH 07/33] pretrain over heedb with pclr --- .github/CODEOWNERS | 1 + configs/pretrain-unsupervised-no-proto.yaml | 52 +++++++ protossl/defines.py | 1 + protossl/models/__init__.py | 1 + protossl/models/_blackbox_contraster.py | 143 ++++++++++++++++++++ protossl/trainer.py | 41 +++++- scripts/pretrain/run-pclr-heedb-pretrain.sh | 42 ++++++ 7 files changed, 280 insertions(+), 1 deletion(-) create mode 100644 configs/pretrain-unsupervised-no-proto.yaml create mode 100644 protossl/models/_blackbox_contraster.py create mode 100755 scripts/pretrain/run-pclr-heedb-pretrain.sh diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index eac1b2b..41986bf 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -34,6 +34,7 @@ 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/_prototype_assigner.py @StevenSong protossl/models/_prototype_classifier.py @StevenSong diff --git a/configs/pretrain-unsupervised-no-proto.yaml b/configs/pretrain-unsupervised-no-proto.yaml new file mode 100644 index 0000000..4a86318 --- /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/protossl/defines.py b/protossl/defines.py index 37b77c6..f6ebc64 100644 --- a/protossl/defines.py +++ b/protossl/defines.py @@ -11,6 +11,7 @@ "project-prototypes-supervised", "compute-embeddings", "train-classifier", + "train-contraster", ] LABEL_T = Literal["binary-multilabel", "multiclass"] 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/trainer.py b/protossl/trainer.py index 5d0e89e..0daa088 100644 --- a/protossl/trainer.py +++ b/protossl/trainer.py @@ -31,6 +31,7 @@ from .models import ( BaseClassifier, BlackboxClassifier, + BlackboxContraster, PrototypeAssigner, PrototypeClassifier, PrototypeContraster, @@ -96,7 +97,7 @@ def setup(self, stage: str): contrastive_pair_mode: CONTRASTIVE_T | None = self.hparams.contrastive_pair_mode # type: ignore data_kwargs: dict = self.hparams.data_kwargs # type: ignore data_wrapper_kwargs: dict = self.hparams.data_wrapper_kwargs # type: ignore - wrap_contrastive = pipeline_stage == "learn-prototypes" + wrap_contrastive = pipeline_stage in {"learn-prototypes", "train-contraster"} if stage == "fit": if not hasattr(self, "train_ds"): @@ -322,6 +323,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 @@ -516,6 +530,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, @@ -1009,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/scripts/pretrain/run-pclr-heedb-pretrain.sh b/scripts/pretrain/run-pclr-heedb-pretrain.sh new file mode 100755 index 0000000..c7b4313 --- /dev/null +++ b/scripts/pretrain/run-pclr-heedb-pretrain.sh @@ -0,0 +1,42 @@ +#!/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/pclr-heedb-%j.out + +set -e + +PRETRAIN_DATASET=/opt/gpudata/ecg/heedb +RUN_DIR=/home/songs1/protossl-ecg-outputs-rebuttal +REPO_ROOT=/home/songs1/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 \ + --config $REPO_ROOT/configs/pretrain-unsupervised-no-proto.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 From 94ece6aaa3c047d24a02860408a7aa67a6ccd0b8 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 23 Jul 2026 23:32:02 -0500 Subject: [PATCH 08/33] run on nvl01 --- configs/pretrain-unsupervised-no-proto.yaml | 2 +- scripts/pretrain/run-pclr-heedb-pretrain.sh | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/configs/pretrain-unsupervised-no-proto.yaml b/configs/pretrain-unsupervised-no-proto.yaml index 4a86318..67fa57a 100644 --- a/configs/pretrain-unsupervised-no-proto.yaml +++ b/configs/pretrain-unsupervised-no-proto.yaml @@ -38,7 +38,7 @@ lr_scheduler: patience: 3 # check early stopping model: class_path: LitModel - init_args: + # init_args: # backbone_type: resnet18 # conv_type: 2D # pretrained_weights: /path/to/weights diff --git a/scripts/pretrain/run-pclr-heedb-pretrain.sh b/scripts/pretrain/run-pclr-heedb-pretrain.sh index c7b4313..af8a5f6 100755 --- a/scripts/pretrain/run-pclr-heedb-pretrain.sh +++ b/scripts/pretrain/run-pclr-heedb-pretrain.sh @@ -4,7 +4,7 @@ #SBATCH --mem-per-gpu=500gb #SBATCH --gpus-per-node=1 #SBATCH --nodes=1 -#SBATCH -w kg35-nvl02 +#SBATCH -w kg35-nvl01 #SBATCH --ntasks-per-node=1 #SBATCH --time=0 #SBATCH --output /home/songs1/slurm-logs/pclr-heedb-%j.out @@ -12,8 +12,8 @@ set -e PRETRAIN_DATASET=/opt/gpudata/ecg/heedb -RUN_DIR=/home/songs1/protossl-ecg-outputs-rebuttal -REPO_ROOT=/home/songs1/ProtoSSL +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}" @@ -34,9 +34,11 @@ 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 From f17b0912b8a8a705180b2568fd0eef72e6fac01f Mon Sep 17 00:00:00 2001 From: StevenSong Date: Fri, 24 Jul 2026 08:52:49 -0500 Subject: [PATCH 09/33] run heedb-150 downstream --- .github/CODEOWNERS | 7 +- scripts/3-run-protossl-heedb-150-pila.sh | 110 ++++++++++++++++ scripts/4-run-supproto-heedb-150-rila.sh | 118 ++++++++++++++++++ ...b-rila.sh => 4-run-supproto-heedb-rila.sh} | 0 scripts/queue-experiments.sh | 78 ++++++------ 5 files changed, 270 insertions(+), 43 deletions(-) create mode 100755 scripts/3-run-protossl-heedb-150-pila.sh create mode 100755 scripts/4-run-supproto-heedb-150-rila.sh rename scripts/{4-run-labsup-proto-heedb-rila.sh => 4-run-supproto-heedb-rila.sh} (100%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 41986bf..2126a9d 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -68,11 +68,8 @@ scripts/0-run-cache-data.sh @StevenSong scripts/1-run-blackbox-direct.sh @StevenSong scripts/2-run-labsup-proto-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/6-run-protossl-heedb-pia.sh @StevenSong 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/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 100% rename from scripts/4-run-labsup-proto-heedb-rila.sh rename to scripts/4-run-supproto-heedb-rila.sh diff --git a/scripts/queue-experiments.sh b/scripts/queue-experiments.sh index a08c60f..5826b7a 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" @@ -45,43 +45,45 @@ for seed in "${SEEDS[@]}"; do cache_id=$(submit_job "$suffix" 0-run-cache-data.sh) 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" 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" 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) - - 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" 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" - - # supproto no-proj - submit_job "$suffix" 4-y-run-labsup-proto-heedb-pila.sh "--dependency=afterok:$cache_id" - - # assignment method - submit_job "$suffix" 6-run-protossl-heedb-pia.sh "--dependency=afterok:$cache_id" - - # 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" - - # start with pretrained encoder - patches_id=$(submit_job "$suffix" 9-0-run-ecgfounder-patches.sh) - echo $patches_id - submit_job "$suffix" 9-1-run-ecgfounder-lap.sh "--dependency=afterok:$patches_id" - submit_job "$suffix" 9-2-run-ecgfounder-clustering.sh "--dependency=afterok:$patches_id" - submit_job "$suffix" 9-3-run-ecgfounder-random.sh "--dependency=afterok:$patches_id" - fi + # 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" 3-run-protossl-heedb-pila.sh "--dependency=afterok:$cache_id" + # submit_job "$suffix" 4-run-supproto-heedb-rila.sh "--dependency=afterok:$cache_id" + 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" 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) + + # 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" 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" + + # # supproto no-proj + # submit_job "$suffix" 4-y-run-labsup-proto-heedb-pila.sh "--dependency=afterok:$cache_id" + + # # assignment method + # submit_job "$suffix" 6-run-protossl-heedb-pia.sh "--dependency=afterok:$cache_id" + + # # 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" + + # # start with pretrained encoder + # patches_id=$(submit_job "$suffix" 9-0-run-ecgfounder-patches.sh) + # echo $patches_id + # submit_job "$suffix" 9-1-run-ecgfounder-lap.sh "--dependency=afterok:$patches_id" + # submit_job "$suffix" 9-2-run-ecgfounder-clustering.sh "--dependency=afterok:$patches_id" + # submit_job "$suffix" 9-3-run-ecgfounder-random.sh "--dependency=afterok:$patches_id" + # fi done done done From d903258738f830aa4a783c77caa84a37a8c3c6f5 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Fri, 24 Jul 2026 13:22:50 -0500 Subject: [PATCH 10/33] pretrain 1050 prototypes --- .../pretrain-supervised-heedb-150-1050.yaml | 55 ++++++++ .../pretrain-unsupervised-heedb-150-1050.yaml | 56 +++++++++ scripts/3-run-protossl-heedb-150-1050-pila.sh | 110 ++++++++++++++++ scripts/4-run-supproto-heedb-150-1050-rila.sh | 118 ++++++++++++++++++ .../run-protossl-heedb-150-1050-pretrain.sh | 75 +++++++++++ .../run-supproto-heedb-150-1050-pretrain.sh | 81 ++++++++++++ scripts/queue-experiments.sh | 10 +- 7 files changed, 501 insertions(+), 4 deletions(-) create mode 100644 configs/pretrain-supervised-heedb-150-1050.yaml create mode 100644 configs/pretrain-unsupervised-heedb-150-1050.yaml create mode 100755 scripts/3-run-protossl-heedb-150-1050-pila.sh create mode 100755 scripts/4-run-supproto-heedb-150-1050-rila.sh create mode 100755 scripts/pretrain/run-protossl-heedb-150-1050-pretrain.sh create mode 100755 scripts/pretrain/run-supproto-heedb-150-1050-pretrain.sh 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-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/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/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/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-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/queue-experiments.sh b/scripts/queue-experiments.sh index 5826b7a..bbc2dc3 100755 --- a/scripts/queue-experiments.sh +++ b/scripts/queue-experiments.sh @@ -42,15 +42,17 @@ for seed in "${SEEDS[@]}"; do for suffix in "${suffixes[@]}"; do echo "SUFFIX: $suffix" - cache_id=$(submit_job "$suffix" 0-run-cache-data.sh) - echo $cache_id + # cache_id=$(submit_job "$suffix" 0-run-cache-data.sh) + # 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" 3-run-protossl-heedb-pila.sh "--dependency=afterok:$cache_id" # submit_job "$suffix" 4-run-supproto-heedb-rila.sh "--dependency=afterok:$cache_id" - 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-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:25565" + submit_job "$suffix" 4-run-supproto-heedb-150-1050-rila.sh "--dependency=afterok:25564" # 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) From d9ea9cada6e34199c4c96071f8e55bebd3f53bfc Mon Sep 17 00:00:00 2001 From: StevenSong Date: Fri, 24 Jul 2026 16:15:24 -0500 Subject: [PATCH 11/33] pclr probe --- .github/CODEOWNERS | 4 +- scripts/5-3-run-pclr-logreg.sh | 39 +++++++++ scripts/ecg-fms/_compute_pclr_embeddings.py | 87 +++++++++++++++++++++ scripts/queue-experiments.sh | 5 +- 4 files changed, 131 insertions(+), 4 deletions(-) create mode 100755 scripts/5-3-run-pclr-logreg.sh create mode 100644 scripts/ecg-fms/_compute_pclr_embeddings.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2126a9d..e8bd53d 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -72,6 +72,7 @@ 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 @@ -81,8 +82,7 @@ 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 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/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/queue-experiments.sh b/scripts/queue-experiments.sh index bbc2dc3..f260ace 100755 --- a/scripts/queue-experiments.sh +++ b/scripts/queue-experiments.sh @@ -51,11 +51,12 @@ for seed in "${SEEDS[@]}"; do # submit_job "$suffix" 4-run-supproto-heedb-rila.sh "--dependency=afterok:$cache_id" # 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:25565" - submit_job "$suffix" 4-run-supproto-heedb-150-1050-rila.sh "--dependency=afterok:25564" + # submit_job "$suffix" 3-run-protossl-heedb-150-1050-pila.sh "--dependency=afterok:25565" + # submit_job "$suffix" 4-run-supproto-heedb-150-1050-rila.sh "--dependency=afterok:25564" # 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 # if [ "$dataset" == "echonext" ]; then # echo "Ablations" From 01d85c4917cfc3f3e99ccaba2e2ddab2b7644668 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 25 Jul 2026 15:51:02 -0500 Subject: [PATCH 12/33] generalize plotting code --- .github/CODEOWNERS | 1 + protossl/plotting.py | 653 +++++++++++++++++++++++++++++++++++++++++++ protossl/trainer.py | 1 - 3 files changed, 654 insertions(+), 1 deletion(-) create mode 100644 protossl/plotting.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index e8bd53d..9fb85a1 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -54,6 +54,7 @@ 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 pyproject.toml @StevenSong README.md @StevenSong @sahilsethi0105 diff --git a/protossl/plotting.py b/protossl/plotting.py new file mode 100644 index 0000000..59ebaaa --- /dev/null +++ b/protossl/plotting.py @@ -0,0 +1,653 @@ +"""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, +) -> None: + """Draw a cutout of a (L, T) segment, one subplot per lead of the layout.""" + 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 + 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/trainer.py b/protossl/trainer.py index 0daa088..95151d7 100644 --- a/protossl/trainer.py +++ b/protossl/trainer.py @@ -580,7 +580,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": From cc7dfed86bafd476b55c003cf938063604dd2fc0 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 25 Jul 2026 15:52:49 -0500 Subject: [PATCH 13/33] prototype recurrence --- results/proto-recur.ipynb | 369 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 369 insertions(+) create mode 100644 results/proto-recur.ipynb diff --git a/results/proto-recur.ipynb b/results/proto-recur.ipynb new file mode 100644 index 0000000..654348c --- /dev/null +++ b/results/proto-recur.ipynb @@ -0,0 +1,369 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "from protossl.defines import PTBXL_TARGETS, CINC_TARGETS, CODE15_TARGETS, HEEDB_TARGETS\n", + "\n", + "from pathlib import Path\n", + "\n", + "run_dir = Path(\"/opt/gpu_working/steven/protossl-ecg-outputs\")\n", + "ptbxl_dir = run_dir / \"experiments-seed42/runs-ptbxl\"\n", + "cinc_dir = run_dir / \"experiments-seed42/runs-cinc\"\n", + "code15_dir = run_dir / \"experiments-seed42/runs-code15\"\n", + "\n", + "PPL = 14" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "(\n", + " list(HEEDB_TARGETS).index(\"RIGHT BUNDLE BRANCH BLOCK\"),\n", + " PTBXL_TARGETS.index(\"CRBBB\"),\n", + " CINC_TARGETS.index(\"CRBBB\"),\n", + " CODE15_TARGETS.index(\"RBBB\"),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "heedb_proj = pd.read_csv(run_dir / \"protossl-heedb/project-prototypes/latest/projection_metadata.csv\")\n", + "\n", + "ptbxl_assign = pd.read_csv(ptbxl_dir / \"protossl-heedb-pila/learn-prototype-assignments/latest/assignment_metadata.csv\")\n", + "ptbxl_proj = pd.read_csv(ptbxl_dir / \"protossl-heedb-pila/project-prototypes-supervised/latest/projection_metadata.csv\")\n", + "ptbxl_ft_proj = pd.read_csv(ptbxl_dir / \"protossl-heedb-pila-ft/project-prototypes-supervised/latest/projection_metadata.csv\")\n", + "\n", + "cinc_assign = pd.read_csv(cinc_dir / \"protossl-heedb-pila/learn-prototype-assignments/latest/assignment_metadata.csv\")\n", + "cinc_proj = pd.read_csv(cinc_dir / \"protossl-heedb-pila/project-prototypes-supervised/latest/projection_metadata.csv\")\n", + "cinc_ft_proj = pd.read_csv(cinc_dir / \"protossl-heedb-pila-ft/project-prototypes-supervised/latest/projection_metadata.csv\")\n", + "\n", + "code15_assign = pd.read_csv(code15_dir / \"protossl-heedb-pila/learn-prototype-assignments/latest/assignment_metadata.csv\")\n", + "code15_proj = pd.read_csv(code15_dir / \"protossl-heedb-pila/project-prototypes-supervised/latest/projection_metadata.csv\")\n", + "code15_ft_proj = pd.read_csv(code15_dir / \"protossl-heedb-pila-ft/project-prototypes-supervised/latest/projection_metadata.csv\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "ptbxl_protos = set(ptbxl_assign.loc[ptbxl_assign[\"label_name\"] == \"CRBBB\", \"prototype_idx\"])\n", + "cinc_protos = set(cinc_assign.loc[cinc_assign[\"label_name\"] == \"CRBBB\", \"prototype_idx\"])\n", + "code15_protos = set(code15_assign.loc[code15_assign[\"label_name\"] == \"RBBB\", \"prototype_idx\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "sets = {\n", + " \"PTB-XL\": ptbxl_protos,\n", + " \"CinC\": cinc_protos,\n", + " \"CODE-15%\": code15_protos,\n", + "}\n", + "names = list(sets)\n", + "universe = sorted(set().union(*sets.values()))\n", + "\n", + "flags = pd.DataFrame(\n", + " {name: [x in s for x in universe] for name, s in sets.items()},\n", + " index=pd.Index(universe, name=\"value\"),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "regions = (\n", + " flags.groupby(names)\n", + " .agg(count=(\"PTB-XL\", \"size\"), members=(flags.columns[0], lambda s: list(s.index)))\n", + " .reset_index()\n", + ")\n", + "regions[\"region\"] = \"\"\n", + "for i in range(len(regions)):\n", + " temp = [n for n in names if regions.loc[i, n]]\n", + " regions.loc[i, \"region\"] = \" & \".join(temp) + (\" only\" if len(temp) == 1 else \"\")\n", + "regions[\"num_regions\"] = regions[names].sum(axis=1)\n", + "\n", + "out = regions.sort_values([\"num_regions\", \"region\"]).reset_index(drop=True)\n", + "n = len(out)\n", + "assert out[\"count\"].sum() == len(ptbxl_protos | cinc_protos | code15_protos)\n", + "out.loc[n, \"count\"] = out[\"count\"].sum()\n", + "out.loc[n, \"region\"] = \"Total\"\n", + "out[\"count\"] = out[\"count\"].astype(int)\n", + "out[[\"region\", \"count\", \"members\"]]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "out[[\"region\", \"count\"]].rename(columns={\"region\": \"Datasets\", \"count\": \"Count\"})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "# !pip install matplotlib-venn\n", + "from matplotlib_venn import venn3\n", + "\n", + "venn3(\n", + " tuple(sets.values()),\n", + " set_labels=tuple(sets.keys()),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "proto_idx = 394" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "heedb_proj_sample = heedb_proj.loc[proto_idx]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "ptbxl_slot = ptbxl_assign[ptbxl_assign[\"prototype_idx\"] == proto_idx]\n", + "assert len(ptbxl_slot) == 1\n", + "ptbxl_proj_sample = ptbxl_proj.loc[ptbxl_slot.iloc[0].name]\n", + "ptbxl_ft_proj_sample = ptbxl_ft_proj.loc[ptbxl_slot.iloc[0].name]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "cinc_slot = cinc_assign[cinc_assign[\"prototype_idx\"] == proto_idx]\n", + "assert len(cinc_slot) == 1\n", + "cinc_proj_sample = cinc_proj.loc[cinc_slot.iloc[0].name]\n", + "cinc_ft_proj_sample = cinc_ft_proj.loc[cinc_slot.iloc[0].name]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "code15_slot = code15_assign[code15_assign[\"prototype_idx\"] == proto_idx]\n", + "assert len(code15_slot) == 1\n", + "code15_proj_sample = code15_proj.loc[code15_slot.iloc[0].name]\n", + "code15_ft_proj_sample = code15_ft_proj.loc[code15_slot.iloc[0].name]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "from protossl.datasets import HeedbECGDataset, PtbxlECGDataset, CincECGDataset, Code15ECGDataset\n", + "from protossl.plotting import plot_ecg\n", + "\n", + "heedb_ds = HeedbECGDataset(dataset_path=\"/opt/gpudata/ecg/heedb\", split=\"train\", sampling_rate=100)\n", + "ptbxl_ds = PtbxlECGDataset(dataset_path=\"/opt/gpudata/ecg/ptb-xl\", split=\"train\", sampling_rate=100)\n", + "cinc_ds = CincECGDataset(dataset_path=\"/opt/gpudata/ecg/cinc-2020\", split=\"train\", sampling_rate=100)\n", + "code15_ds = Code15ECGDataset(dataset_path=\"/opt/gpudata/ecg/code15\", split=\"train\", sampling_rate=100)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "heedb_fig = plot_ecg(dataset=heedb_ds, sample_id=int(heedb_proj_sample[\"ecg_id\"]), chunk_idx=int(heedb_proj_sample[\"chunk_idx\"]))\n", + "heedb_fig.savefig(\"figs/heedb-rbbb.png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "ptbxl_fig = plot_ecg(dataset=ptbxl_ds, sample_id=int(ptbxl_proj_sample[\"ecg_id\"]), chunk_idx=int(ptbxl_proj_sample[\"chunk_idx\"]))\n", + "ptbxl_fig.savefig(\"figs/ptbxl-rbbb.png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "ptbxl_ft_fig = plot_ecg(dataset=ptbxl_ds, sample_id=int(ptbxl_ft_proj_sample[\"ecg_id\"]), chunk_idx=int(ptbxl_ft_proj_sample[\"chunk_idx\"]))\n", + "ptbxl_ft_fig.savefig(\"figs/ptbxl-ft-rbbb.png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "cinc_fig = plot_ecg(dataset=cinc_ds, sample_id=int(cinc_proj_sample[\"ecg_id\"]), chunk_idx=int(cinc_proj_sample[\"chunk_idx\"]))\n", + "cinc_fig.savefig(\"figs/cinc-rbbb.png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "cinc_ft_fig = plot_ecg(dataset=cinc_ds, sample_id=int(cinc_ft_proj_sample[\"ecg_id\"]), chunk_idx=int(cinc_ft_proj_sample[\"chunk_idx\"]))\n", + "cinc_ft_fig.savefig(\"figs/cinc-ft-rbbb.png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "code15_fig = plot_ecg(dataset=code15_ds, sample_id=int(code15_proj_sample[\"ecg_id\"]), chunk_idx=int(code15_proj_sample[\"chunk_idx\"]))\n", + "code15_fig.savefig(\"figs/code15-rbbb.png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20", + "metadata": {}, + "outputs": [], + "source": [ + "code15_ft_fig = plot_ecg(dataset=code15_ds, sample_id=int(code15_ft_proj_sample[\"ecg_id\"]), chunk_idx=int(code15_ft_proj_sample[\"chunk_idx\"]))\n", + "code15_ft_fig.savefig(\"figs/code15-ft-rbbb.png\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "import torch" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22", + "metadata": {}, + "outputs": [], + "source": [ + "sd = torch.load(\"../scripts/ecgfounder-checkpoint/12_lead_ECGFounder.pth\", map_location=\"cpu\", weights_only=False)[\"state_dict\"]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "sd[\"dense.weight\"].shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "sd[\"dense.bias\"].shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "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 +} From 17093388431461b2e733b3abe535a820dc63662b Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 25 Jul 2026 16:36:15 -0500 Subject: [PATCH 14/33] wip recur multiple --- results/proto-recur-multilabel.ipynb | 641 +++++++++++++++++++++++++++ 1 file changed, 641 insertions(+) create mode 100644 results/proto-recur-multilabel.ipynb diff --git a/results/proto-recur-multilabel.ipynb b/results/proto-recur-multilabel.ipynb new file mode 100644 index 0000000..eb0805b --- /dev/null +++ b/results/proto-recur-multilabel.ipynb @@ -0,0 +1,641 @@ +{ + "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 HEEDB waveform 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 generalizes the single-label (RBBB) version in `proto-recur.ipynb` to\n", + "an arbitrary set of 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 as projected by every run that uses it\n", + " (HEEDB pretraining, plus each downstream dataset's probe and fine-tuned models).\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 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", + " CincECGDataset,\n", + " Code15ECGDataset,\n", + " HeedbECGDataset,\n", + " PtbxlECGDataset,\n", + ")\n", + "from protossl.defines import (\n", + " CINC_TARGETS,\n", + " CODE15_TARGETS,\n", + " HEEDB_TARGETS,\n", + " PTBXL_TARGETS,\n", + ")\n", + "from protossl.plotting import plot_ecg, resolve_index\n", + "\n", + "# runs\n", + "RUN_DIR = Path(\"/opt/gpu_working/steven/protossl-ecg-outputs\")\n", + "PRETRAIN_DIR = RUN_DIR / \"protossl-heedb\" # HEEDB pretraining, source of all prototypes\n", + "EXPERIMENT_DIR = RUN_DIR / \"experiments-seed42\" # downstream runs, one dir per dataset\n", + "MODEL = \"protossl-heedb-pila\" # linear probe run, \"-ft\" suffix is the fine-tuned run\n", + "\n", + "# data\n", + "DATA_DIR = Path(\"/opt/gpudata/ecg\")\n", + "SAMPLING_RATE = 100 # models were trained at 100 Hz\n", + "SPLIT = \"train\" # prototypes are projected onto the train split (see protossl.trainer)\n", + "\n", + "# model, see configs/target-guided-14ppl.yaml\n", + "PPL = 14 # prototypes per label (n_prototypes_per_label)\n", + "\n", + "FIG_DIR = Path(\"figs\")\n", + "FIG_DIR.mkdir(parents=True, exist_ok=True)" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "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": "3", + "metadata": {}, + "outputs": [], + "source": [ + "PRETRAIN = {\n", + " \"slug\": \"heedb\",\n", + " \"dataset_cls\": HeedbECGDataset,\n", + " \"dataset_path\": DATA_DIR / \"heedb\",\n", + " \"targets\": list(HEEDB_TARGETS), # a dict of label -> statement codes\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": "4", + "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": "5", + "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": "6", + "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.\n", + " \"\"\"\n", + " return pd.read_csv(path).rename(columns={\"ecg_id\": \"sample_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\n", + " / f\"{MODEL}-ft\"\n", + " / \"project-prototypes-supervised/latest/projection_metadata.csv\"\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "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": "8", + "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": "9", + "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": "10", + "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": "11", + "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\" if size == 1 else \"\"),\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": "12", + "metadata": {}, + "outputs": [], + "source": [ + "LABEL = \"RBBB\" # label to look at in detail below\n", + "\n", + "add_total(region_table(PROTOTYPES[LABEL]))[[\"region\", \"count\", \"prototypes\"]]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "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\"] = [\n", + " len(set().union(*sets.values())) for sets in PROTOTYPES.values()\n", + "]\n", + "overlap.rename_axis(\"Datasets\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "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": "15", + "metadata": {}, + "outputs": [], + "source": [ + "# !pip install matplotlib-venn\n", + "from matplotlib_venn import venn3\n", + "\n", + "assert len(DATASETS) == 3, \"venn3 draws exactly three sets\"\n", + "\n", + "n_cols = 3\n", + "n_rows = -(-len(PROTOTYPES) // n_cols) # ceiling division\n", + "fig, axs = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows))\n", + "for ax, (label, sets) in zip(axs.flat, PROTOTYPES.items()):\n", + " venn3([sets[name] for name in DATASETS], set_labels=tuple(DATASETS), ax=ax)\n", + " ax.set_title(label)\n", + "for ax in axs.flat[len(PROTOTYPES) :]: # blank out any unused panels\n", + " ax.axis(\"off\")\n", + "fig.tight_layout()\n", + "fig.savefig(FIG_DIR / \"proto-recur-venn.png\", dpi=300, bbox_inches=\"tight\")" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "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": "17", + "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": "18", + "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": "19", + "metadata": {}, + "outputs": [], + "source": [ + "def prototype_sources(proto_idx: int) -> 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", + " 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": "20", + "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 = True,\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).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", + " )\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": "21", + "metadata": {}, + "outputs": [], + "source": [ + "# prefer a prototype every dataset agreed on, then any recurrent one, then any at all\n", + "candidates = recurrent_prototypes(LABEL, min_datasets=len(DATASETS))\n", + "if not candidates:\n", + " candidates = recurrent_prototypes(LABEL)\n", + "if not candidates:\n", + " candidates = sorted(set().union(*PROTOTYPES[LABEL].values()))\n", + "\n", + "PROTO_IDX = candidates[0] # or any prototype index of interest\n", + "\n", + "prototype_sources(PROTO_IDX)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22", + "metadata": {}, + "outputs": [], + "source": [ + "figs = plot_prototype(LABEL, PROTO_IDX)\n", + "list(figs)" + ] + }, + { + "cell_type": "markdown", + "id": "23", + "metadata": {}, + "source": [ + "## Optional: HEEDB labels of the prototype's source ECG\n", + "\n", + "The pretraining is unsupervised, so a prototype has no HEEDB label of its own. What\n", + "it does have is a source ECG, and that ECG's HEEDB statements are a useful sanity\n", + "check on whether the prototype really encodes the concept the downstream datasets\n", + "assigned it to. Set `USE_HEEDB_150=1` before importing `protossl` if the pretraining\n", + "run used the 150-label HEEDB set, otherwise these names come from the 20-label set." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "def heedb_labels(sample_id: int, split: str = SPLIT) -> list[str]:\n", + " \"\"\"HEEDB statements carried by a source ECG.\"\"\"\n", + " ds = get_dataset(\"HEEDB\", split=split)\n", + " idx = resolve_index(ds, sample_id=sample_id)\n", + " return [\n", + " name\n", + " for name, present in zip(PRETRAIN[\"targets\"], ds.labels[idx].tolist())\n", + " if present\n", + " ]\n", + "\n", + "\n", + "expected = LABELS[LABEL][\"HEEDB\"]\n", + "ecg_id = int(heedb_proj.loc[PROTO_IDX, \"sample_id\"])\n", + "statements = heedb_labels(ecg_id)\n", + "print(f\"prototype {PROTO_IDX} came from HEEDB ECG {ecg_id}\")\n", + "seen = \"present\" if expected in statements else \"absent\"\n", + "print(f\"expected statement {expected}: {seen}\")\n", + "statements" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 46c47f68a5b8d95b14024f9278e88174f2961b6b Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 25 Jul 2026 19:24:39 -0500 Subject: [PATCH 15/33] cleanup multilabel proto-recur --- results/proto-recur-multilabel.ipynb | 172 +++++-------- results/proto-recur.ipynb | 369 --------------------------- 2 files changed, 69 insertions(+), 472 deletions(-) delete mode 100644 results/proto-recur.ipynb diff --git a/results/proto-recur-multilabel.ipynb b/results/proto-recur-multilabel.ipynb index eb0805b..abf28e8 100644 --- a/results/proto-recur-multilabel.ipynb +++ b/results/proto-recur-multilabel.ipynb @@ -10,17 +10,17 @@ "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 HEEDB waveform is doing\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 generalizes the single-label (RBBB) version in `proto-recur.ipynb` to\n", - "an arbitrary set of labels and produces:\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 as projected by every run that uses it\n", - " (HEEDB pretraining, plus each downstream dataset's probe and fine-tuned models).\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." @@ -40,32 +40,30 @@ "import pandas as pd\n", "\n", "from protossl.datasets import (\n", + " HeedbECGDataset,\n", " CincECGDataset,\n", " Code15ECGDataset,\n", - " HeedbECGDataset,\n", " PtbxlECGDataset,\n", ")\n", "from protossl.defines import (\n", + " HEEDB_TARGETS,\n", " CINC_TARGETS,\n", " CODE15_TARGETS,\n", - " HEEDB_TARGETS,\n", " PTBXL_TARGETS,\n", ")\n", - "from protossl.plotting import plot_ecg, resolve_index\n", + "from protossl.plotting import plot_ecg\n", + "\n", + "SEED = 42\n", "\n", - "# runs\n", "RUN_DIR = Path(\"/opt/gpu_working/steven/protossl-ecg-outputs\")\n", - "PRETRAIN_DIR = RUN_DIR / \"protossl-heedb\" # HEEDB pretraining, source of all prototypes\n", - "EXPERIMENT_DIR = RUN_DIR / \"experiments-seed42\" # downstream runs, one dir per dataset\n", - "MODEL = \"protossl-heedb-pila\" # linear probe run, \"-ft\" suffix is the fine-tuned run\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\n", "DATA_DIR = Path(\"/opt/gpudata/ecg\")\n", - "SAMPLING_RATE = 100 # models were trained at 100 Hz\n", - "SPLIT = \"train\" # prototypes are projected onto the train split (see protossl.trainer)\n", - "\n", - "# model, see configs/target-guided-14ppl.yaml\n", - "PPL = 14 # prototypes per label (n_prototypes_per_label)\n", + "SAMPLING_RATE = 100\n", + "SPLIT = \"train\"\n", "\n", "FIG_DIR = Path(\"figs\")\n", "FIG_DIR.mkdir(parents=True, exist_ok=True)" @@ -105,7 +103,7 @@ " \"slug\": \"heedb\",\n", " \"dataset_cls\": HeedbECGDataset,\n", " \"dataset_path\": DATA_DIR / \"heedb\",\n", - " \"targets\": list(HEEDB_TARGETS), # a dict of label -> statement codes\n", + " \"targets\": list(HEEDB_TARGETS),\n", "}\n", "\n", "DATASETS = {\n", @@ -226,9 +224,9 @@ "\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.\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\"})\n", + " return pd.read_csv(path).rename(columns={\"ecg_id\": \"sample_id\", \"patient_id\": \"source_id\"})\n", "\n", "\n", "heedb_proj = read_metadata(\n", @@ -244,9 +242,7 @@ " run / MODEL / \"project-prototypes-supervised/latest/projection_metadata.csv\"\n", " )\n", " cfg[\"ft_proj\"] = read_metadata(\n", - " run\n", - " / f\"{MODEL}-ft\"\n", - " / \"project-prototypes-supervised/latest/projection_metadata.csv\"\n", + " run / f\"{MODEL}-ft\" / \"project-prototypes-supervised/latest/projection_metadata.csv\"\n", " )" ] }, @@ -344,7 +340,7 @@ " members = shared.difference(*others) if others else shared\n", " rows.append(\n", " {\n", - " \"region\": \" & \".join(combo) + (\" only\" if size == 1 else \"\"),\n", + " \"region\": \" & \".join(combo) + \" only\",\n", " \"n_datasets\": size,\n", " \"count\": len(members),\n", " \"prototypes\": sorted(members),\n", @@ -367,18 +363,6 @@ "id": "12", "metadata": {}, "outputs": [], - "source": [ - "LABEL = \"RBBB\" # label to look at in detail below\n", - "\n", - "add_total(region_table(PROTOTYPES[LABEL]))[[\"region\", \"count\", \"prototypes\"]]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "13", - "metadata": {}, - "outputs": [], "source": [ "# every label at once: rows are Venn regions, columns are labels\n", "overlap = pd.DataFrame(\n", @@ -396,7 +380,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -415,7 +399,7 @@ { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "14", "metadata": {}, "outputs": [], "source": [ @@ -433,12 +417,12 @@ "for ax in axs.flat[len(PROTOTYPES) :]: # blank out any unused panels\n", " ax.axis(\"off\")\n", "fig.tight_layout()\n", - "fig.savefig(FIG_DIR / \"proto-recur-venn.png\", dpi=300, bbox_inches=\"tight\")" + "fig.savefig(FIG_DIR / f\"proto-recur-venn-seed{SEED}.png\", dpi=300, bbox_inches=\"tight\")" ] }, { "cell_type": "markdown", - "id": "16", + "id": "15", "metadata": {}, "source": [ "## Recurrent prototypes\n", @@ -450,7 +434,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -469,7 +453,7 @@ }, { "cell_type": "markdown", - "id": "18", + "id": "17", "metadata": {}, "source": [ "## Plot a prototype's projections\n", @@ -482,26 +466,27 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "18", "metadata": {}, "outputs": [], "source": [ - "def prototype_sources(proto_idx: int) -> pd.DataFrame:\n", + "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", - " 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", + " 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))" @@ -510,7 +495,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "19", "metadata": {}, "outputs": [], "source": [ @@ -534,11 +519,12 @@ " label: str,\n", " proto_idx: int,\n", " split: str = SPLIT,\n", - " save: bool = True,\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).iterrows():\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", @@ -557,83 +543,63 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "20", "metadata": {}, "outputs": [], "source": [ - "# prefer a prototype every dataset agreed on, then any recurrent one, then any at all\n", - "candidates = recurrent_prototypes(LABEL, min_datasets=len(DATASETS))\n", - "if not candidates:\n", - " candidates = recurrent_prototypes(LABEL)\n", - "if not candidates:\n", - " candidates = sorted(set().union(*PROTOTYPES[LABEL].values()))\n", - "\n", - "PROTO_IDX = candidates[0] # or any prototype index of interest\n", - "\n", - "prototype_sources(PROTO_IDX)" + "# 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": "22", + "id": "21", "metadata": {}, "outputs": [], "source": [ - "figs = plot_prototype(LABEL, PROTO_IDX)\n", - "list(figs)" + "# let's look at a highly-recurrent prototype\n", + "PROTO_IDX = 394\n", + "prototype_sources(PROTO_IDX)" ] }, { - "cell_type": "markdown", - "id": "23", + "cell_type": "code", + "execution_count": null, + "id": "22", "metadata": {}, + "outputs": [], "source": [ - "## Optional: HEEDB labels of the prototype's source ECG\n", - "\n", - "The pretraining is unsupervised, so a prototype has no HEEDB label of its own. What\n", - "it does have is a source ECG, and that ECG's HEEDB statements are a useful sanity\n", - "check on whether the prototype really encodes the concept the downstream datasets\n", - "assigned it to. Set `USE_HEEDB_150=1` before importing `protossl` if the pretraining\n", - "run used the 150-label HEEDB set, otherwise these names come from the 20-label set." + "figs = plot_prototype(\"RBBB\", PROTO_IDX, pretrain_only=True)" ] }, { "cell_type": "code", "execution_count": null, - "id": "24", + "id": "23", "metadata": {}, "outputs": [], - "source": [ - "def heedb_labels(sample_id: int, split: str = SPLIT) -> list[str]:\n", - " \"\"\"HEEDB statements carried by a source ECG.\"\"\"\n", - " ds = get_dataset(\"HEEDB\", split=split)\n", - " idx = resolve_index(ds, sample_id=sample_id)\n", - " return [\n", - " name\n", - " for name, present in zip(PRETRAIN[\"targets\"], ds.labels[idx].tolist())\n", - " if present\n", - " ]\n", - "\n", - "\n", - "expected = LABELS[LABEL][\"HEEDB\"]\n", - "ecg_id = int(heedb_proj.loc[PROTO_IDX, \"sample_id\"])\n", - "statements = heedb_labels(ecg_id)\n", - "print(f\"prototype {PROTO_IDX} came from HEEDB ECG {ecg_id}\")\n", - "seen = \"present\" if expected in statements else \"absent\"\n", - "print(f\"expected statement {expected}: {seen}\")\n", - "statements" - ] + "source": [] } ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "protossl", "language": "python", "name": "python3" }, "language_info": { - "name": "python" + "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, diff --git a/results/proto-recur.ipynb b/results/proto-recur.ipynb deleted file mode 100644 index 654348c..0000000 --- a/results/proto-recur.ipynb +++ /dev/null @@ -1,369 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "import pandas as pd\n", - "from protossl.defines import PTBXL_TARGETS, CINC_TARGETS, CODE15_TARGETS, HEEDB_TARGETS\n", - "\n", - "from pathlib import Path\n", - "\n", - "run_dir = Path(\"/opt/gpu_working/steven/protossl-ecg-outputs\")\n", - "ptbxl_dir = run_dir / \"experiments-seed42/runs-ptbxl\"\n", - "cinc_dir = run_dir / \"experiments-seed42/runs-cinc\"\n", - "code15_dir = run_dir / \"experiments-seed42/runs-code15\"\n", - "\n", - "PPL = 14" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "(\n", - " list(HEEDB_TARGETS).index(\"RIGHT BUNDLE BRANCH BLOCK\"),\n", - " PTBXL_TARGETS.index(\"CRBBB\"),\n", - " CINC_TARGETS.index(\"CRBBB\"),\n", - " CODE15_TARGETS.index(\"RBBB\"),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "heedb_proj = pd.read_csv(run_dir / \"protossl-heedb/project-prototypes/latest/projection_metadata.csv\")\n", - "\n", - "ptbxl_assign = pd.read_csv(ptbxl_dir / \"protossl-heedb-pila/learn-prototype-assignments/latest/assignment_metadata.csv\")\n", - "ptbxl_proj = pd.read_csv(ptbxl_dir / \"protossl-heedb-pila/project-prototypes-supervised/latest/projection_metadata.csv\")\n", - "ptbxl_ft_proj = pd.read_csv(ptbxl_dir / \"protossl-heedb-pila-ft/project-prototypes-supervised/latest/projection_metadata.csv\")\n", - "\n", - "cinc_assign = pd.read_csv(cinc_dir / \"protossl-heedb-pila/learn-prototype-assignments/latest/assignment_metadata.csv\")\n", - "cinc_proj = pd.read_csv(cinc_dir / \"protossl-heedb-pila/project-prototypes-supervised/latest/projection_metadata.csv\")\n", - "cinc_ft_proj = pd.read_csv(cinc_dir / \"protossl-heedb-pila-ft/project-prototypes-supervised/latest/projection_metadata.csv\")\n", - "\n", - "code15_assign = pd.read_csv(code15_dir / \"protossl-heedb-pila/learn-prototype-assignments/latest/assignment_metadata.csv\")\n", - "code15_proj = pd.read_csv(code15_dir / \"protossl-heedb-pila/project-prototypes-supervised/latest/projection_metadata.csv\")\n", - "code15_ft_proj = pd.read_csv(code15_dir / \"protossl-heedb-pila-ft/project-prototypes-supervised/latest/projection_metadata.csv\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [ - "ptbxl_protos = set(ptbxl_assign.loc[ptbxl_assign[\"label_name\"] == \"CRBBB\", \"prototype_idx\"])\n", - "cinc_protos = set(cinc_assign.loc[cinc_assign[\"label_name\"] == \"CRBBB\", \"prototype_idx\"])\n", - "code15_protos = set(code15_assign.loc[code15_assign[\"label_name\"] == \"RBBB\", \"prototype_idx\"])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [ - "sets = {\n", - " \"PTB-XL\": ptbxl_protos,\n", - " \"CinC\": cinc_protos,\n", - " \"CODE-15%\": code15_protos,\n", - "}\n", - "names = list(sets)\n", - "universe = sorted(set().union(*sets.values()))\n", - "\n", - "flags = pd.DataFrame(\n", - " {name: [x in s for x in universe] for name, s in sets.items()},\n", - " index=pd.Index(universe, name=\"value\"),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [ - "regions = (\n", - " flags.groupby(names)\n", - " .agg(count=(\"PTB-XL\", \"size\"), members=(flags.columns[0], lambda s: list(s.index)))\n", - " .reset_index()\n", - ")\n", - "regions[\"region\"] = \"\"\n", - "for i in range(len(regions)):\n", - " temp = [n for n in names if regions.loc[i, n]]\n", - " regions.loc[i, \"region\"] = \" & \".join(temp) + (\" only\" if len(temp) == 1 else \"\")\n", - "regions[\"num_regions\"] = regions[names].sum(axis=1)\n", - "\n", - "out = regions.sort_values([\"num_regions\", \"region\"]).reset_index(drop=True)\n", - "n = len(out)\n", - "assert out[\"count\"].sum() == len(ptbxl_protos | cinc_protos | code15_protos)\n", - "out.loc[n, \"count\"] = out[\"count\"].sum()\n", - "out.loc[n, \"region\"] = \"Total\"\n", - "out[\"count\"] = out[\"count\"].astype(int)\n", - "out[[\"region\", \"count\", \"members\"]]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [ - "out[[\"region\", \"count\"]].rename(columns={\"region\": \"Datasets\", \"count\": \"Count\"})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [ - "# !pip install matplotlib-venn\n", - "from matplotlib_venn import venn3\n", - "\n", - "venn3(\n", - " tuple(sets.values()),\n", - " set_labels=tuple(sets.keys()),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8", - "metadata": {}, - "outputs": [], - "source": [ - "proto_idx = 394" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9", - "metadata": {}, - "outputs": [], - "source": [ - "heedb_proj_sample = heedb_proj.loc[proto_idx]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "10", - "metadata": {}, - "outputs": [], - "source": [ - "ptbxl_slot = ptbxl_assign[ptbxl_assign[\"prototype_idx\"] == proto_idx]\n", - "assert len(ptbxl_slot) == 1\n", - "ptbxl_proj_sample = ptbxl_proj.loc[ptbxl_slot.iloc[0].name]\n", - "ptbxl_ft_proj_sample = ptbxl_ft_proj.loc[ptbxl_slot.iloc[0].name]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "11", - "metadata": {}, - "outputs": [], - "source": [ - "cinc_slot = cinc_assign[cinc_assign[\"prototype_idx\"] == proto_idx]\n", - "assert len(cinc_slot) == 1\n", - "cinc_proj_sample = cinc_proj.loc[cinc_slot.iloc[0].name]\n", - "cinc_ft_proj_sample = cinc_ft_proj.loc[cinc_slot.iloc[0].name]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "12", - "metadata": {}, - "outputs": [], - "source": [ - "code15_slot = code15_assign[code15_assign[\"prototype_idx\"] == proto_idx]\n", - "assert len(code15_slot) == 1\n", - "code15_proj_sample = code15_proj.loc[code15_slot.iloc[0].name]\n", - "code15_ft_proj_sample = code15_ft_proj.loc[code15_slot.iloc[0].name]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "13", - "metadata": {}, - "outputs": [], - "source": [ - "from protossl.datasets import HeedbECGDataset, PtbxlECGDataset, CincECGDataset, Code15ECGDataset\n", - "from protossl.plotting import plot_ecg\n", - "\n", - "heedb_ds = HeedbECGDataset(dataset_path=\"/opt/gpudata/ecg/heedb\", split=\"train\", sampling_rate=100)\n", - "ptbxl_ds = PtbxlECGDataset(dataset_path=\"/opt/gpudata/ecg/ptb-xl\", split=\"train\", sampling_rate=100)\n", - "cinc_ds = CincECGDataset(dataset_path=\"/opt/gpudata/ecg/cinc-2020\", split=\"train\", sampling_rate=100)\n", - "code15_ds = Code15ECGDataset(dataset_path=\"/opt/gpudata/ecg/code15\", split=\"train\", sampling_rate=100)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "14", - "metadata": {}, - "outputs": [], - "source": [ - "heedb_fig = plot_ecg(dataset=heedb_ds, sample_id=int(heedb_proj_sample[\"ecg_id\"]), chunk_idx=int(heedb_proj_sample[\"chunk_idx\"]))\n", - "heedb_fig.savefig(\"figs/heedb-rbbb.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "15", - "metadata": {}, - "outputs": [], - "source": [ - "ptbxl_fig = plot_ecg(dataset=ptbxl_ds, sample_id=int(ptbxl_proj_sample[\"ecg_id\"]), chunk_idx=int(ptbxl_proj_sample[\"chunk_idx\"]))\n", - "ptbxl_fig.savefig(\"figs/ptbxl-rbbb.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "16", - "metadata": {}, - "outputs": [], - "source": [ - "ptbxl_ft_fig = plot_ecg(dataset=ptbxl_ds, sample_id=int(ptbxl_ft_proj_sample[\"ecg_id\"]), chunk_idx=int(ptbxl_ft_proj_sample[\"chunk_idx\"]))\n", - "ptbxl_ft_fig.savefig(\"figs/ptbxl-ft-rbbb.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "17", - "metadata": {}, - "outputs": [], - "source": [ - "cinc_fig = plot_ecg(dataset=cinc_ds, sample_id=int(cinc_proj_sample[\"ecg_id\"]), chunk_idx=int(cinc_proj_sample[\"chunk_idx\"]))\n", - "cinc_fig.savefig(\"figs/cinc-rbbb.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "18", - "metadata": {}, - "outputs": [], - "source": [ - "cinc_ft_fig = plot_ecg(dataset=cinc_ds, sample_id=int(cinc_ft_proj_sample[\"ecg_id\"]), chunk_idx=int(cinc_ft_proj_sample[\"chunk_idx\"]))\n", - "cinc_ft_fig.savefig(\"figs/cinc-ft-rbbb.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "19", - "metadata": {}, - "outputs": [], - "source": [ - "code15_fig = plot_ecg(dataset=code15_ds, sample_id=int(code15_proj_sample[\"ecg_id\"]), chunk_idx=int(code15_proj_sample[\"chunk_idx\"]))\n", - "code15_fig.savefig(\"figs/code15-rbbb.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "20", - "metadata": {}, - "outputs": [], - "source": [ - "code15_ft_fig = plot_ecg(dataset=code15_ds, sample_id=int(code15_ft_proj_sample[\"ecg_id\"]), chunk_idx=int(code15_ft_proj_sample[\"chunk_idx\"]))\n", - "code15_ft_fig.savefig(\"figs/code15-ft-rbbb.png\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "21", - "metadata": {}, - "outputs": [], - "source": [ - "import torch" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "22", - "metadata": {}, - "outputs": [], - "source": [ - "sd = torch.load(\"../scripts/ecgfounder-checkpoint/12_lead_ECGFounder.pth\", map_location=\"cpu\", weights_only=False)[\"state_dict\"]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "23", - "metadata": {}, - "outputs": [], - "source": [ - "sd[\"dense.weight\"].shape" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "24", - "metadata": {}, - "outputs": [], - "source": [ - "sd[\"dense.bias\"].shape" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "25", - "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 -} From 3dbc879c2d5b65875d9b7fe9500b577280dba4fc Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 25 Jul 2026 19:26:23 -0500 Subject: [PATCH 16/33] plot rebuttal exp results --- results/ecg-results-rebuttal.ipynb | 706 +++++++++++++++++++++++++++++ 1 file changed, 706 insertions(+) create mode 100644 results/ecg-results-rebuttal.ipynb diff --git a/results/ecg-results-rebuttal.ipynb b/results/ecg-results-rebuttal.ipynb new file mode 100644 index 0000000..76c6c57 --- /dev/null +++ b/results/ecg-results-rebuttal.ipynb @@ -0,0 +1,706 @@ +{ + "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\"),\n", + " ###\n", + " \"SupProto HEEDB-150 (7500-proto-bank)\": (\"tab:orange\", \"supproto-heedb-150-rila\"),\n", + " \"SupProto HEEDB-150 (7500-proto-bank) (FT)\": (\"tab:pink\", \"supproto-heedb-150-rila-ft\"),\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank)\": (\"tab:blue\", \"protossl-heedb-150-pila\"),\n", + " \"ProtoSSL HEEDB-150 (7500-proto-bank) (FT)\": (\"tab:cyan\", \"protossl-heedb-150-pila-ft\"),\n", + " ###\n", + " \"SupProto HEEDB-150 (1050-proto-bank)\": (\"tab:orange\", \"supproto-heedb-150-1050-rila\"),\n", + " \"SupProto HEEDB-150 (1050-proto-bank) (FT)\": (\"tab:pink\", \"supproto-heedb-150-1050-rila-ft\"),\n", + " \"ProtoSSL HEEDB-150 (1050-proto-bank)\": (\"tab:blue\", \"protossl-heedb-150-1050-pila\"),\n", + " \"ProtoSSL HEEDB-150 (1050-proto-bank) (FT)\": (\"tab:cyan\", \"protossl-heedb-150-1050-pila-ft\"),\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": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## Make Tables - TODO cleanup below" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "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": "13", + "metadata": {}, + "source": [ + "### Full Data Scale, Primary Models" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "temp = pivoted.loc[\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", + " [\"EchoNext\", \"MIMIC-IV-ECG\", \"ZZU pECG\", \"PTB-XL\", \"CinC Georgia\", \"CODE-15%\"]\n", + "].T\n", + "temp = temp.reset_index().drop_duplicates(\"Dataset\", keep=\"first\").drop(columns=[\"Train Size\"]).reset_index(drop=True)\n", + "temp = pd.concat([pd.DataFrame({\"Domain\": [\"Out-of-Domain\"]*3+[\"In-Domain\"]*3}), temp], axis=1)\n", + "temp = temp.set_index([\"Domain\", \"Dataset\"])\n", + "temp.columns = pd.MultiIndex.from_arrays([\n", + " [\"Tuned\", \"Tuned\", \"Probed\", \"Probed\", \"\"],\n", + " [\"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", + "])\n", + "\n", + "print(temp.to_latex())" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "### All Data Scales, Primary Models" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", + " [\"EchoNext\", \"MIMIC-IV-ECG\", \"ZZU pECG\", \"PTB-XL\", \"CinC Georgia\", \"CODE-15%\"]\n", + " ].T.to_latex()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "17", + "metadata": {}, + "source": [ + "### Blackbox Baselines" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB\", \"ECGFounder\", \"ST-MEM\", \"Blackbox Direct\"],\n", + " [\"EchoNext\", \"MIMIC-IV-ECG\", \"ZZU pECG\", \"PTB-XL\", \"CinC Georgia\", \"CODE-15%\"]\n", + " ].T.to_latex()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": {}, + "source": [ + "### Ablations" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "#### Num Prototypes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "ppl7 = pivoted.loc[\n", + " [\"ProtoSSL HEEDB (7ppl) (FT)\", \"SupProto HEEDB (7ppl) (FT)\", \"ProtoSSL HEEDB (7ppl)\", \"SupProto HEEDB (7ppl)\", \"SupProto Direct (7ppl)\"],\n", + " [\"EchoNext\"]\n", + " ].T.copy().rename(columns=dict(zip(\n", + " [\"ProtoSSL HEEDB (7ppl) (FT)\", \"SupProto HEEDB (7ppl) (FT)\", \"ProtoSSL HEEDB (7ppl)\", \"SupProto HEEDB (7ppl)\", \"SupProto Direct (7ppl)\"],\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"]\n", + " )))\n", + "\n", + "ppl14 = pivoted.loc[\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", + " [\"EchoNext\"]\n", + " ].T.copy()\n", + "\n", + "ppl28 = pivoted.loc[\n", + " [\"ProtoSSL HEEDB (28ppl) (FT)\", \"SupProto HEEDB (28ppl) (FT)\", \"ProtoSSL HEEDB (28ppl)\", \"SupProto HEEDB (28ppl)\", \"SupProto Direct (28ppl)\"],\n", + " [\"EchoNext\"]\n", + " ].T.copy().rename(columns=dict(zip(\n", + " [\"ProtoSSL HEEDB (28ppl) (FT)\", \"SupProto HEEDB (28ppl) (FT)\", \"ProtoSSL HEEDB (28ppl)\", \"SupProto HEEDB (28ppl)\", \"SupProto Direct (28ppl)\"],\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"]\n", + " )))\n", + "\n", + "ppl7[\"Prototypes Per Label\"] = 7\n", + "ppl14[\"Prototypes Per Label\"] = 14\n", + "ppl28[\"Prototypes Per Label\"] = 28\n", + "\n", + "ppl7.set_index(\"Prototypes Per Label\", append=True, inplace=True)\n", + "ppl14.set_index(\"Prototypes Per Label\", append=True, inplace=True)\n", + "ppl28.set_index(\"Prototypes Per Label\", append=True, inplace=True)\n", + "\n", + "ppl_df = pd.concat([ppl7, ppl14, ppl28]).reorder_levels([0, 2, 1]).reset_index(level=0, drop=True)\n", + "print(ppl_df.to_latex())" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "#### SupProto NoProj" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "pila = pivoted.loc[\n", + " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (PILA) (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB (PILA)\", \"SupProto HEEDB\"],\n", + " [\"EchoNext\"]\n", + " ].T.copy()\n", + "pila.columns = pd.MultiIndex.from_tuples([\n", + " (\"Tuned\", \"ProtoSSL HEEDB\"),\n", + " (\"Tuned\", \"SupProto HEEDB (NoProj)\"),\n", + " (\"Tuned\", \"SupProto HEEDB\"),\n", + " (\"Probed\", \"ProtoSSL HEEDB\"),\n", + " (\"Probed\", \"SupProto HEEDB (NoProj)\"),\n", + " (\"Probed\", \"SupProto HEEDB\"),\n", + "])\n", + "print(pila.reset_index(level=0, drop=True).to_latex())" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "#### LAP vs ProtoPool Assignment" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB (FT)\", \"ProtoSSL HEEDB (PIA) (FT)\", \"ProtoSSL HEEDB\", \"ProtoSSL HEEDB (PIA)\"],\n", + " [\"EchoNext\"]\n", + " ].T.rename(columns={\n", + " \"ProtoSSL HEEDB\": \"ProtoSSL HEEDB (LAP)\",\n", + " \"ProtoSSL HEEDB (PIA)\": \"ProtoSSL HEEDB (Pool)\",\n", + " }).to_latex()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "#### No Assignment (PIT & PIP)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB (83ppl)\", \"ProtoSSL HEEDB (PIT)\", \"ProtoSSL HEEDB (PIP)\"],\n", + " [\"EchoNext\"]\n", + " ].T.rename(columns={\n", + " \"ProtoSSL HEEDB (83ppl)\": \"ProtoSSL HEEDB (LAP) (83PPL)\",\n", + " }).to_latex()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "### Prototypes from FM" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "print(\n", + " pivoted.loc[\n", + " [\"ProtoSSL HEEDB\", \"ECGFounder (SK-OT)\", \"ECGFounder (LAP)\", \"ECGFounder (Rand)\", \"ECGFounder\"],\n", + " [\"EchoNext\"]\n", + " ].T.to_latex()\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "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 +} From 74fa72519f5db7960cdd0d840fe1b885164682bf Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 25 Jul 2026 20:12:08 -0500 Subject: [PATCH 17/33] queue latent prototype probe --- .github/CODEOWNERS | 3 +- ...b-pit.sh => 7-1-run-protossl-heedb-pit.sh} | 0 ...b-pip.sh => 7-2-run-protossl-heedb-pip.sh} | 0 scripts/7-3-run-protossl-heedb-latent.sh | 43 +++++++++++++++++++ scripts/queue-experiments.sh | 19 ++++---- 5 files changed, 55 insertions(+), 10 deletions(-) rename scripts/{7-run-protossl-heedb-pit.sh => 7-1-run-protossl-heedb-pit.sh} (100%) rename scripts/{8-run-protossl-heedb-pip.sh => 7-2-run-protossl-heedb-pip.sh} (100%) create mode 100755 scripts/7-3-run-protossl-heedb-latent.sh diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 9fb85a1..fcf9469 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -76,8 +76,7 @@ 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 diff --git a/scripts/7-run-protossl-heedb-pit.sh b/scripts/7-1-run-protossl-heedb-pit.sh similarity index 100% rename from scripts/7-run-protossl-heedb-pit.sh rename to scripts/7-1-run-protossl-heedb-pit.sh diff --git a/scripts/8-run-protossl-heedb-pip.sh b/scripts/7-2-run-protossl-heedb-pip.sh similarity index 100% rename from scripts/8-run-protossl-heedb-pip.sh rename to scripts/7-2-run-protossl-heedb-pip.sh 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/queue-experiments.sh b/scripts/queue-experiments.sh index f260ace..b77fd38 100755 --- a/scripts/queue-experiments.sh +++ b/scripts/queue-experiments.sh @@ -2,7 +2,8 @@ set -e -declare -a DATASETS=("echonext" "ptbxl" "cinc" "mimic" "zzu" "code15") +declare -a DATASETS=("echonext") +# declare -a DATASETS=("echonext" "ptbxl" "cinc" "mimic" "zzu" "code15") declare -a SEEDS=(42 67 70 73 99) declare -A DATASET_DIRS=( ["echonext"]="echonext" @@ -51,14 +52,14 @@ for seed in "${SEEDS[@]}"; do # submit_job "$suffix" 4-run-supproto-heedb-rila.sh "--dependency=afterok:$cache_id" # 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:25565" - # submit_job "$suffix" 4-run-supproto-heedb-150-1050-rila.sh "--dependency=afterok:25564" + # 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" # 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 + # submit_job "$suffix" 5-3-run-pclr-logreg.sh "--dependency=afterok:$cache_id" - # if [ "$dataset" == "echonext" ]; then + if [ "$dataset" == "echonext" ]; then # echo "Ablations" # # prototypes per label # submit_job "$suffix" 2-z-run-labsup-proto-direct-7ppl.sh "--dependency=afterok:$cache_id" @@ -77,8 +78,10 @@ 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" + submit_job "$suffix" 7-3-run-protossl-heedb-latent.sh # # start with pretrained encoder # patches_id=$(submit_job "$suffix" 9-0-run-ecgfounder-patches.sh) @@ -86,7 +89,7 @@ for seed in "${SEEDS[@]}"; do # submit_job "$suffix" 9-1-run-ecgfounder-lap.sh "--dependency=afterok:$patches_id" # submit_job "$suffix" 9-2-run-ecgfounder-clustering.sh "--dependency=afterok:$patches_id" # submit_job "$suffix" 9-3-run-ecgfounder-random.sh "--dependency=afterok:$patches_id" - # fi + fi done done done From a5a12110bf85af1122a000441e9fd6613ff62670 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Tue, 28 Jul 2026 08:07:59 -0500 Subject: [PATCH 18/33] rebuttal results --- results/ecg-results-rebuttal.ipynb | 277 +++++++++++---------------- results/proto-recur-multilabel.ipynb | 96 ++++++++-- 2 files changed, 193 insertions(+), 180 deletions(-) diff --git a/results/ecg-results-rebuttal.ipynb b/results/ecg-results-rebuttal.ipynb index 76c6c57..a55d88d 100644 --- a/results/ecg-results-rebuttal.ipynb +++ b/results/ecg-results-rebuttal.ipynb @@ -107,17 +107,22 @@ " \"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\"),\n", + " \"PCLR-HEEDB\": (\"tab:olive\", \"pclr-logreg\"), # new\n", " ###\n", - " \"SupProto HEEDB-150 (7500-proto-bank)\": (\"tab:orange\", \"supproto-heedb-150-rila\"),\n", - " \"SupProto HEEDB-150 (7500-proto-bank) (FT)\": (\"tab:pink\", \"supproto-heedb-150-rila-ft\"),\n", - " \"ProtoSSL HEEDB-150 (7500-proto-bank)\": (\"tab:blue\", \"protossl-heedb-150-pila\"),\n", - " \"ProtoSSL HEEDB-150 (7500-proto-bank) (FT)\": (\"tab:cyan\", \"protossl-heedb-150-pila-ft\"),\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\"),\n", - " \"SupProto HEEDB-150 (1050-proto-bank) (FT)\": (\"tab:pink\", \"supproto-heedb-150-1050-rila-ft\"),\n", - " \"ProtoSSL HEEDB-150 (1050-proto-bank)\": (\"tab:blue\", \"protossl-heedb-150-1050-pila\"),\n", - " \"ProtoSSL HEEDB-150 (1050-proto-bank) (FT)\": (\"tab:cyan\", \"protossl-heedb-150-1050-pila-ft\"),\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", @@ -382,11 +387,27 @@ ] }, { - "cell_type": "markdown", + "cell_type": "code", + "execution_count": null, "id": "11", "metadata": {}, + "outputs": [], "source": [ - "## Make Tables - TODO cleanup below" + "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])" ] }, { @@ -395,6 +416,56 @@ "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", @@ -439,86 +510,36 @@ }, { "cell_type": "markdown", - "id": "13", - "metadata": {}, - "source": [ - "### Full Data Scale, Primary Models" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "14", - "metadata": {}, - "outputs": [], - "source": [ - "temp = pivoted.loc[\n", - " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", - " [\"EchoNext\", \"MIMIC-IV-ECG\", \"ZZU pECG\", \"PTB-XL\", \"CinC Georgia\", \"CODE-15%\"]\n", - "].T\n", - "temp = temp.reset_index().drop_duplicates(\"Dataset\", keep=\"first\").drop(columns=[\"Train Size\"]).reset_index(drop=True)\n", - "temp = pd.concat([pd.DataFrame({\"Domain\": [\"Out-of-Domain\"]*3+[\"In-Domain\"]*3}), temp], axis=1)\n", - "temp = temp.set_index([\"Domain\", \"Dataset\"])\n", - "temp.columns = pd.MultiIndex.from_arrays([\n", - " [\"Tuned\", \"Tuned\", \"Probed\", \"Probed\", \"\"],\n", - " [\"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", - "])\n", - "\n", - "print(temp.to_latex())" - ] - }, - { - "cell_type": "markdown", - "id": "15", - "metadata": {}, - "source": [ - "### All Data Scales, Primary Models" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "16", + "id": "17", "metadata": {}, - "outputs": [], "source": [ - "print(\n", - " pivoted.loc[\n", - " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", - " [\"EchoNext\", \"MIMIC-IV-ECG\", \"ZZU pECG\", \"PTB-XL\", \"CinC Georgia\", \"CODE-15%\"]\n", - " ].T.to_latex()\n", - ")" + "### HEEDB 150 Labels (new)" ] }, { "cell_type": "markdown", - "id": "17", + "id": "18", "metadata": {}, "source": [ - "### Blackbox Baselines" + "#### 1050 Pretrained Prototypes (7 PPL during pretraining for SupProto HEEDB)" ] }, { "cell_type": "code", "execution_count": null, - "id": "18", - "metadata": {}, - "outputs": [], - "source": [ - "print(\n", - " pivoted.loc[\n", - " [\"ProtoSSL HEEDB\", \"ECGFounder\", \"ST-MEM\", \"Blackbox Direct\"],\n", - " [\"EchoNext\", \"MIMIC-IV-ECG\", \"ZZU pECG\", \"PTB-XL\", \"CinC Georgia\", \"CODE-15%\"]\n", - " ].T.to_latex()\n", - ")" - ] - }, - { - "cell_type": "markdown", "id": "19", "metadata": {}, + "outputs": [], "source": [ - "### Ablations" + "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())" ] }, { @@ -526,7 +547,7 @@ "id": "20", "metadata": {}, "source": [ - "#### Num Prototypes" + "#### 7500 Pretrained Prototypes (50 PPL during pretraining for SupProto HEEDB)" ] }, { @@ -536,37 +557,15 @@ "metadata": {}, "outputs": [], "source": [ - "ppl7 = pivoted.loc[\n", - " [\"ProtoSSL HEEDB (7ppl) (FT)\", \"SupProto HEEDB (7ppl) (FT)\", \"ProtoSSL HEEDB (7ppl)\", \"SupProto HEEDB (7ppl)\", \"SupProto Direct (7ppl)\"],\n", - " [\"EchoNext\"]\n", - " ].T.copy().rename(columns=dict(zip(\n", - " [\"ProtoSSL HEEDB (7ppl) (FT)\", \"SupProto HEEDB (7ppl) (FT)\", \"ProtoSSL HEEDB (7ppl)\", \"SupProto HEEDB (7ppl)\", \"SupProto Direct (7ppl)\"],\n", - " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"]\n", - " )))\n", - "\n", - "ppl14 = pivoted.loc[\n", - " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB\", \"SupProto Direct\"],\n", - " [\"EchoNext\"]\n", - " ].T.copy()\n", - "\n", - "ppl28 = pivoted.loc[\n", - " [\"ProtoSSL HEEDB (28ppl) (FT)\", \"SupProto HEEDB (28ppl) (FT)\", \"ProtoSSL HEEDB (28ppl)\", \"SupProto HEEDB (28ppl)\", \"SupProto Direct (28ppl)\"],\n", + "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 (28ppl) (FT)\", \"SupProto HEEDB (28ppl) (FT)\", \"ProtoSSL HEEDB (28ppl)\", \"SupProto HEEDB (28ppl)\", \"SupProto Direct (28ppl)\"],\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", - "ppl7[\"Prototypes Per Label\"] = 7\n", - "ppl14[\"Prototypes Per Label\"] = 14\n", - "ppl28[\"Prototypes Per Label\"] = 28\n", - "\n", - "ppl7.set_index(\"Prototypes Per Label\", append=True, inplace=True)\n", - "ppl14.set_index(\"Prototypes Per Label\", append=True, inplace=True)\n", - "ppl28.set_index(\"Prototypes Per Label\", append=True, inplace=True)\n", - "\n", - "ppl_df = pd.concat([ppl7, ppl14, ppl28]).reorder_levels([0, 2, 1]).reset_index(level=0, drop=True)\n", - "print(ppl_df.to_latex())" + "print(n7500.to_markdown())" ] }, { @@ -574,7 +573,7 @@ "id": "22", "metadata": {}, "source": [ - "#### SupProto NoProj" + "### Blackbox Baselines - including PCLR (new)" ] }, { @@ -583,100 +582,44 @@ "id": "23", "metadata": {}, "outputs": [], - "source": [ - "pila = pivoted.loc[\n", - " [\"ProtoSSL HEEDB (FT)\", \"SupProto HEEDB (PILA) (FT)\", \"SupProto HEEDB (FT)\", \"ProtoSSL HEEDB\", \"SupProto HEEDB (PILA)\", \"SupProto HEEDB\"],\n", - " [\"EchoNext\"]\n", - " ].T.copy()\n", - "pila.columns = pd.MultiIndex.from_tuples([\n", - " (\"Tuned\", \"ProtoSSL HEEDB\"),\n", - " (\"Tuned\", \"SupProto HEEDB (NoProj)\"),\n", - " (\"Tuned\", \"SupProto HEEDB\"),\n", - " (\"Probed\", \"ProtoSSL HEEDB\"),\n", - " (\"Probed\", \"SupProto HEEDB (NoProj)\"),\n", - " (\"Probed\", \"SupProto HEEDB\"),\n", - "])\n", - "print(pila.reset_index(level=0, drop=True).to_latex())" - ] - }, - { - "cell_type": "markdown", - "id": "24", - "metadata": {}, - "source": [ - "#### LAP vs ProtoPool Assignment" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "25", - "metadata": {}, - "outputs": [], "source": [ "print(\n", " pivoted.loc[\n", - " [\"ProtoSSL HEEDB (FT)\", \"ProtoSSL HEEDB (PIA) (FT)\", \"ProtoSSL HEEDB\", \"ProtoSSL HEEDB (PIA)\"],\n", + " [\"PCLR-HEEDB\", \"ProtoSSL HEEDB\", \"ECGFounder\", \"ST-MEM\", \"Blackbox Direct\"],\n", " [\"EchoNext\"]\n", - " ].T.rename(columns={\n", - " \"ProtoSSL HEEDB\": \"ProtoSSL HEEDB (LAP)\",\n", - " \"ProtoSSL HEEDB (PIA)\": \"ProtoSSL HEEDB (Pool)\",\n", - " }).to_latex()\n", + " ].T.to_markdown()\n", ")" ] }, { "cell_type": "markdown", - "id": "26", + "id": "24", "metadata": {}, "source": [ - "#### No Assignment (PIT & PIP)" + "### No Assignment (PIT & PIP) - and No Projection (Latent) (new)" ] }, { "cell_type": "code", "execution_count": null, - "id": "27", + "id": "25", "metadata": {}, "outputs": [], "source": [ "print(\n", " pivoted.loc[\n", - " [\"ProtoSSL HEEDB (83ppl)\", \"ProtoSSL HEEDB (PIT)\", \"ProtoSSL HEEDB (PIP)\"],\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_latex()\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "28", - "metadata": {}, - "source": [ - "### Prototypes from FM" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "29", - "metadata": {}, - "outputs": [], - "source": [ - "print(\n", - " pivoted.loc[\n", - " [\"ProtoSSL HEEDB\", \"ECGFounder (SK-OT)\", \"ECGFounder (LAP)\", \"ECGFounder (Rand)\", \"ECGFounder\"],\n", - " [\"EchoNext\"]\n", - " ].T.to_latex()\n", + " }).to_markdown()\n", ")" ] }, { "cell_type": "code", "execution_count": null, - "id": "30", + "id": "26", "metadata": {}, "outputs": [], "source": [] diff --git a/results/proto-recur-multilabel.ipynb b/results/proto-recur-multilabel.ipynb index abf28e8..bcc4bc5 100644 --- a/results/proto-recur-multilabel.ipynb +++ b/results/proto-recur-multilabel.ipynb @@ -363,6 +363,74 @@ "id": "12", "metadata": {}, "outputs": [], + "source": [ + "region_table(PROTOTYPES[\"Sinus Brady\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "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": "14", + "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": "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" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], "source": [ "# every label at once: rows are Venn regions, columns are labels\n", "overlap = pd.DataFrame(\n", @@ -371,16 +439,18 @@ " for label, sets in PROTOTYPES.items()\n", " }\n", ")\n", - "overlap.loc[\"Total\"] = [\n", + "overlap.loc[\"Total Unique\"] = [\n", " len(set().union(*sets.values())) for sets in PROTOTYPES.values()\n", "]\n", - "overlap.rename_axis(\"Datasets\")" + "overlap = overlap.rename_axis(\"Datasets\")\n", + "display(overlap)\n", + "print(overlap.to_markdown())" ] }, { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "17", "metadata": {}, "outputs": [], "source": [ @@ -399,7 +469,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14", + "id": "18", "metadata": {}, "outputs": [], "source": [ @@ -422,7 +492,7 @@ }, { "cell_type": "markdown", - "id": "15", + "id": "19", "metadata": {}, "source": [ "## Recurrent prototypes\n", @@ -434,7 +504,7 @@ { "cell_type": "code", "execution_count": null, - "id": "16", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -453,7 +523,7 @@ }, { "cell_type": "markdown", - "id": "17", + "id": "21", "metadata": {}, "source": [ "## Plot a prototype's projections\n", @@ -466,7 +536,7 @@ { "cell_type": "code", "execution_count": null, - "id": "18", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -495,7 +565,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -543,7 +613,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -555,7 +625,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "25", "metadata": {}, "outputs": [], "source": [ @@ -567,7 +637,7 @@ { "cell_type": "code", "execution_count": null, - "id": "22", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -577,7 +647,7 @@ { "cell_type": "code", "execution_count": null, - "id": "23", + "id": "27", "metadata": {}, "outputs": [], "source": [] From 304a88129add97f4b86bfdbe0d5db7d0a3209f66 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Mon, 17 Aug 2026 08:31:05 -0500 Subject: [PATCH 19/33] final rebuttal result notebook and koleo ablation --- results/ecg-results-rebuttal.ipynb | 64 +++++++++++++++- .../run-protossl-heedb-pretrain-no-koleo.sh | 73 +++++++++++++++++++ .../run-protossl-heedb-pretrain-w-attn.sh | 2 +- scripts/queue-experiments.sh | 13 ++-- 4 files changed, 144 insertions(+), 8 deletions(-) create mode 100755 scripts/pretrain/run-protossl-heedb-pretrain-no-koleo.sh diff --git a/results/ecg-results-rebuttal.ipynb b/results/ecg-results-rebuttal.ipynb index a55d88d..5c733dd 100644 --- a/results/ecg-results-rebuttal.ipynb +++ b/results/ecg-results-rebuttal.ipynb @@ -616,10 +616,72 @@ ")" ] }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "### ZZU full scale more prototypes in ProtoSSL pretraining" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "26", + "id": "27", + "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": "28", + "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": "29", + "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": "30", "metadata": {}, "outputs": [], "source": [] 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-protossl-heedb-pretrain-w-attn.sh b/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh index fda90f2..4841c16 100755 --- a/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh +++ b/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh @@ -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/queue-experiments.sh b/scripts/queue-experiments.sh index b77fd38..e947cd2 100755 --- a/scripts/queue-experiments.sh +++ b/scripts/queue-experiments.sh @@ -2,7 +2,7 @@ set -e -declare -a DATASETS=("echonext") +declare -a DATASETS=("zzu") # declare -a DATASETS=("echonext" "ptbxl" "cinc" "mimic" "zzu" "code15") declare -a SEEDS=(42 67 70 73 99) declare -A DATASET_DIRS=( @@ -18,7 +18,8 @@ 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") +declare -a SUFFIXES_zzu=("") +# declare -a SUFFIXES_zzu=("" "-4k" "-2k" "-1k" "-512" "-256") source _submit_job.sh @@ -50,7 +51,7 @@ for seed in "${SEEDS[@]}"; do # submit_job "$suffix" 2-run-labsup-proto-direct.sh "--dependency=afterok:$cache_id" # submit_job "$suffix" 3-run-protossl-heedb-pila.sh "--dependency=afterok:$cache_id" # submit_job "$suffix" 4-run-supproto-heedb-rila.sh "--dependency=afterok:$cache_id" - # submit_job "$suffix" 3-run-protossl-heedb-150-pila.sh "--dependency=afterok:$cache_id" + submit_job "$suffix" 3-run-protossl-heedb-150-pila.sh # 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" @@ -59,7 +60,7 @@ for seed in "${SEEDS[@]}"; do # 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" == "echonext" ]; then + # if [ "$dataset" == "echonext" ]; then # echo "Ablations" # # prototypes per label # submit_job "$suffix" 2-z-run-labsup-proto-direct-7ppl.sh "--dependency=afterok:$cache_id" @@ -81,7 +82,7 @@ for seed in "${SEEDS[@]}"; do # 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" - submit_job "$suffix" 7-3-run-protossl-heedb-latent.sh + # submit_job "$suffix" 7-3-run-protossl-heedb-latent.sh # # start with pretrained encoder # patches_id=$(submit_job "$suffix" 9-0-run-ecgfounder-patches.sh) @@ -89,7 +90,7 @@ for seed in "${SEEDS[@]}"; do # submit_job "$suffix" 9-1-run-ecgfounder-lap.sh "--dependency=afterok:$patches_id" # submit_job "$suffix" 9-2-run-ecgfounder-clustering.sh "--dependency=afterok:$patches_id" # submit_job "$suffix" 9-3-run-ecgfounder-random.sh "--dependency=afterok:$patches_id" - fi + # fi done done done From e9f8f50e04b0285b7883d15fbe5d84ea04065686 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 27 Aug 2026 16:21:24 -0500 Subject: [PATCH 20/33] protoecgnet over heedb for large user study --- .github/CODEOWNERS | 3 + .gitignore | 2 +- .../configs}/1D-global.yaml | 6 +- .../configs}/2D-global.yaml | 6 +- .../configs}/2D-partial.yaml | 6 +- large-user-study/configs/fusion.yaml | 52 ++ .../configs/ppl-sweep}/wandb-sweep.yaml | 0 .../confirm-new-heedb-splits.ipynb | 227 ++++++ large-user-study/muse-accuracy.ipynb | 761 ++++++++++++++++++ .../scripts/_eval_probs_bootstrapped.py | 142 ++++ .../run-protoecgnet-heedb-1D-global.sh | 73 ++ .../run-protoecgnet-heedb-2D-global.sh | 73 ++ .../run-protoecgnet-heedb-2D-partial.sh | 73 ++ .../scripts/run-protoecgnet-heedb-fusion.sh | 61 ++ protossl/datasets/_heedb_dataset.py | 97 ++- protossl/defines.py | 1 + protossl/models/_protoecgnet.py | 465 +++++++++++ protossl/protoecgnet_trainer.py | 688 ++++++++++++++++ protossl/trainer.py | 1 + requirements.txt | 1 + scripts/pretrain/_cache_heedb.py | 1 + 21 files changed, 2703 insertions(+), 36 deletions(-) rename {configs/multi-arm-heedb => large-user-study/configs}/1D-global.yaml (93%) rename {configs/multi-arm-heedb => large-user-study/configs}/2D-global.yaml (92%) rename {configs/multi-arm-heedb => large-user-study/configs}/2D-partial.yaml (93%) create mode 100644 large-user-study/configs/fusion.yaml rename {configs/multi-arm-heedb => large-user-study/configs/ppl-sweep}/wandb-sweep.yaml (100%) create mode 100644 large-user-study/confirm-new-heedb-splits.ipynb create mode 100644 large-user-study/muse-accuracy.ipynb create mode 100644 large-user-study/scripts/_eval_probs_bootstrapped.py create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-fusion.sh create mode 100644 protossl/models/_protoecgnet.py create mode 100644 protossl/protoecgnet_trainer.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index fcf9469..d0683b4 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -36,6 +36,7 @@ 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 @@ -56,6 +57,7 @@ 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 @@ -95,3 +97,4 @@ 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 dcdba4c..59ef6a8 100644 --- a/.gitignore +++ b/.gitignore @@ -219,4 +219,4 @@ OLD/ slurm-logs/ results/figs/ ecgfounder-checkpoint/ -configs/multi-arm-heedb/wandb +large-user-study/configs/ppl-sweep/wandb diff --git a/configs/multi-arm-heedb/1D-global.yaml b/large-user-study/configs/1D-global.yaml similarity index 93% rename from configs/multi-arm-heedb/1D-global.yaml rename to large-user-study/configs/1D-global.yaml index bdd024f..c4dfc42 100644 --- a/configs/multi-arm-heedb/1D-global.yaml +++ b/large-user-study/configs/1D-global.yaml @@ -40,10 +40,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 +63,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 92% rename from configs/multi-arm-heedb/2D-global.yaml rename to large-user-study/configs/2D-global.yaml index a15df8d..d53562b 100644 --- a/configs/multi-arm-heedb/2D-global.yaml +++ b/large-user-study/configs/2D-global.yaml @@ -40,10 +40,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: 256 # pretrained_weights: /path/to/weights data: class_path: LitData @@ -54,4 +54,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 93% rename from configs/multi-arm-heedb/2D-partial.yaml rename to large-user-study/configs/2D-partial.yaml index a9dacd1..97f18aa 100644 --- a/configs/multi-arm-heedb/2D-partial.yaml +++ b/large-user-study/configs/2D-partial.yaml @@ -40,12 +40,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 +64,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..52ed72e --- /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 # affects 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/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/confirm-new-heedb-splits.ipynb b/large-user-study/confirm-new-heedb-splits.ipynb new file mode 100644 index 0000000..0669ef8 --- /dev/null +++ b/large-user-study/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/muse-accuracy.ipynb b/large-user-study/muse-accuracy.ipynb new file mode 100644 index 0000000..97ab010 --- /dev/null +++ b/large-user-study/muse-accuracy.ipynb @@ -0,0 +1,761 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 29, + "id": "308c1648", + "metadata": {}, + "outputs": [], + "source": [ + "from protossl.datasets import HeedbECGDataset\n", + "from protossl.defines import HEEDB_TARGETS\n", + "\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\"" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "68d44cd8", + "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:04<00:00, 2279302.11it/s]\n", + "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 3200781.88it/s]\n", + "Converting code string to labels: 100%|██████████| 368191/368191 [00:03<00:00, 104052.12it/s]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 368191 ECGs and 11607259 annotations, 368191 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": [ + "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:04<00:00, 2147329.70it/s]\n", + "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 3175092.66it/s]\n", + "Converting code string to labels: 100%|██████████| 368191/368191 [00:01<00:00, 242206.99it/s]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 368191 ECGs and 11607259 annotations, 368191 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": [ + "Creating mgb fname to code mapping: 100%|██████████| 10471531/10471531 [00:04<00:00, 2375692.40it/s]\n", + "Creating emory fname to code mapping: 100%|██████████| 974172/974172 [00:00<00:00, 2824799.44it/s]\n", + "Converting code string to labels: 100%|██████████| 368191/368191 [00:05<00:00, 64363.89it/s] \n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 368191 ECGs and 11445703 annotations, 368191 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", + ")\n", + "\n", + "ds_muse = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"test\",\n", + " sampling_rate=100,\n", + " label_src=\"original_muse\",\n", + ")\n", + "\n", + "ds_muse_v24 = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"test\",\n", + " sampling_rate=100,\n", + " label_src=\"new_muse\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "55a31c9e", + "metadata": {}, + "outputs": [], + "source": [ + "y_true = ds_phys.labels.numpy()\n", + "y_pred_muse = ds_muse.labels.numpy()\n", + "y_pred_v24 = ds_muse_v24.labels.numpy()" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a84dca44", + "metadata": {}, + "outputs": [], + "source": [ + "def compute_pred(_y_true, _y_prob):\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 (_y_prob > optimal_threshold).astype(int)" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "id": "be9ee773", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ANTERIOR INFARCT\n", + "supproto: 0.797\n", + "\n", + "ATRIAL FIBRILLATION\n", + "supproto: 0.962\n", + "\n", + "ATRIAL FLUTTER\n", + "supproto: 0.947\n", + "\n", + "ATRIAL-PACED RHYTHM\n", + "supproto: 0.988\n", + "\n", + "INCOMPLETE RIGHT BUNDLE BRANCH BLOCK\n", + "supproto: 0.967\n", + "\n", + "INFERIOR INFARCT\n", + "supproto: 0.834\n", + "\n", + "LATERAL INFARCT\n", + "supproto: 0.945\n", + "\n", + "LEFT BUNDLE BRANCH BLOCK\n", + "supproto: 0.988\n", + "\n", + "NORMAL SINUS RHYTHM\n", + "supproto: 0.980\n", + "\n", + "PREMATURE ATRIAL COMPLEXES\n", + "supproto: 0.880\n", + "\n", + "PREMATURE VENTRICULAR COMPLEXES\n", + "supproto: 0.875\n", + "\n", + "RIGHT AXIS DEVIATION\n", + "supproto: 0.987\n", + "\n", + "RIGHT BUNDLE BRANCH BLOCK\n", + "supproto: 0.992\n", + "\n", + "SINUS BRADYCARDIA\n", + "supproto: 0.994\n", + "\n", + "SINUS RHYTHM\n", + "supproto: 0.885\n", + "\n", + "SINUS TACHYCARDIA\n", + "supproto: 0.994\n", + "\n", + "VENTRICULAR TACHYCARDIA\n", + "supproto: 0.895\n", + "\n", + "VENTRICULAR-PACED RHYTHM\n", + "supproto: 0.991\n", + "\n", + "WITH 1ST DEGREE AV BLOCK\n", + "supproto: 0.988\n", + "\n", + "WITH SINUS ARRHYTHMIA\n", + "supproto: 0.943\n", + "\n" + ] + } + ], + "source": [ + "y_prob_supproto = np.load(\"/opt/gpu_working/steven/protossl-ecg-outputs/supproto-heedb/probs.npy\")\n", + "y_prob_protossl = np.load(\"/opt/gpu_working/steven/protossl-ecg-outputs/protossl-heedb/probs.npy\")\n", + "\n", + "y_pred_supproto = np.zeros_like(y_pred_muse)\n", + "y_pred_protossl = np.zeros_like(y_pred_muse)\n", + "results = []\n", + "for i, label in enumerate(HEEDB_TARGETS):\n", + " print(label)\n", + " y_pred_supproto[:, i] = compute_pred(y_true[:, i], y_prob_supproto[:, i])\n", + " y_pred_protossl[:, i] = compute_pred(y_true[:, i], y_prob_protossl[:, i])\n", + " label_results = {\n", + " \"label\": label,\n", + " # \"supproto\": balanced_accuracy_score(y_true[:, i], y_pred_supproto[:, i]),\n", + " # \"protossl\": balanced_accuracy_score(y_true[:, i], y_pred_protossl[:, i]),\n", + " # \"old_muse\": balanced_accuracy_score(y_true[:, i], y_pred_muse[:, i]),\n", + " # \"new_muse\": balanced_accuracy_score(y_true[:, i], y_pred_v24[:, i]),\n", + " \"supproto\": roc_auc_score(y_true[:, i], y_prob_supproto[:, i]),\n", + " # \"protossl\": roc_auc_score(y_true[:, i], y_prob_protossl[:, i]),\n", + " }\n", + " results.append(label_results)\n", + " print(f\"supproto: {label_results['supproto']:0.3f}\")\n", + " # print(f\"protossl: {label_results['protossl']:0.3f}\")\n", + " # print(f\"old_muse: {label_results['old_muse']:0.3f}\")\n", + " # print(f\"new_muse: {label_results['new_muse']:0.3f}\")\n", + " print()" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "id": "1e01daec", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " label supproto\n", + " ANTERIOR INFARCT 0.797\n", + " ATRIAL FIBRILLATION 0.962\n", + " ATRIAL FLUTTER 0.947\n", + " ATRIAL-PACED RHYTHM 0.988\n", + "INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.967\n", + " INFERIOR INFARCT 0.834\n", + " LATERAL INFARCT 0.945\n", + " LEFT BUNDLE BRANCH BLOCK 0.988\n", + " NORMAL SINUS RHYTHM 0.980\n", + " PREMATURE ATRIAL COMPLEXES 0.880\n", + " PREMATURE VENTRICULAR COMPLEXES 0.875\n", + " RIGHT AXIS DEVIATION 0.987\n", + " RIGHT BUNDLE BRANCH BLOCK 0.992\n", + " SINUS BRADYCARDIA 0.994\n", + " SINUS RHYTHM 0.885\n", + " SINUS TACHYCARDIA 0.994\n", + " VENTRICULAR TACHYCARDIA 0.895\n", + " VENTRICULAR-PACED RHYTHM 0.991\n", + " WITH 1ST DEGREE AV BLOCK 0.988\n", + " WITH SINUS ARRHYTHMIA 0.943\n" + ] + } + ], + "source": [ + "print(pd.DataFrame(results).to_string(float_format=\"%0.3f\", index=False))" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "id": "56f18902", + "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", + "
labelsupproto
0ANTERIOR INFARCT0.797
1ATRIAL FIBRILLATION0.962
2ATRIAL FLUTTER0.947
3ATRIAL-PACED RHYTHM0.988
4INCOMPLETE RIGHT BUNDLE BRANCH BLOCK0.967
5INFERIOR INFARCT0.834
6LATERAL INFARCT0.945
7LEFT BUNDLE BRANCH BLOCK0.988
8NORMAL SINUS RHYTHM0.980
9PREMATURE ATRIAL COMPLEXES0.880
10PREMATURE VENTRICULAR COMPLEXES0.875
11RIGHT AXIS DEVIATION0.987
12RIGHT BUNDLE BRANCH BLOCK0.992
13SINUS BRADYCARDIA0.994
14SINUS RHYTHM0.885
15SINUS TACHYCARDIA0.994
16VENTRICULAR TACHYCARDIA0.895
17VENTRICULAR-PACED RHYTHM0.991
18WITH 1ST DEGREE AV BLOCK0.988
19WITH SINUS ARRHYTHMIA0.943
\n", + "
" + ], + "text/plain": [ + " label supproto\n", + "0 ANTERIOR INFARCT 0.797\n", + "1 ATRIAL FIBRILLATION 0.962\n", + "2 ATRIAL FLUTTER 0.947\n", + "3 ATRIAL-PACED RHYTHM 0.988\n", + "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.967\n", + "5 INFERIOR INFARCT 0.834\n", + "6 LATERAL INFARCT 0.945\n", + "7 LEFT BUNDLE BRANCH BLOCK 0.988\n", + "8 NORMAL SINUS RHYTHM 0.980\n", + "9 PREMATURE ATRIAL COMPLEXES 0.880\n", + "10 PREMATURE VENTRICULAR COMPLEXES 0.875\n", + "11 RIGHT AXIS DEVIATION 0.987\n", + "12 RIGHT BUNDLE BRANCH BLOCK 0.992\n", + "13 SINUS BRADYCARDIA 0.994\n", + "14 SINUS RHYTHM 0.885\n", + "15 SINUS TACHYCARDIA 0.994\n", + "16 VENTRICULAR TACHYCARDIA 0.895\n", + "17 VENTRICULAR-PACED RHYTHM 0.991\n", + "18 WITH 1ST DEGREE AV BLOCK 0.988\n", + "19 WITH SINUS ARRHYTHMIA 0.943" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "with pd.option_context(\"display.precision\", 3):\n", + " display(pd.DataFrame(results))" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "id": "c2c33568", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "supproto 0.941648\n", + "dtype: float64" + ] + }, + "execution_count": 34, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pd.DataFrame(results).set_index(\"label\").mean(axis=0)" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "d42b88e1", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor(37)" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_phys.labels[:, 0].sum()" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "7e46f74f", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor(0)" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_muse.labels[:, 0].sum()" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "fe2e5a2f", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor(21096)" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_muse_v24.labels[:, 0].sum()" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "a97c7c50", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([368191, 20])" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_phys.labels.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "4f2bef08", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "=================make_heedb_labels=================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Converting code string to labels: 100%|██████████| 8145485/8145485 [01:33<00:00, 87111.97it/s] \n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 8145485 ECGs and 11607259 annotations, 8145484 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" + ] + } + ], + "source": [ + "ds_train_phys = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"train\",\n", + " sampling_rate=100,\n", + " label_src=\"original_physician\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "b1021ed2", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "=================make_heedb_labels=================\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Converting code string to labels: 100%|██████████| 556578/556578 [00:01<00:00, 512100.29it/s]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Of 556578 ECGs and 11607259 annotations, 556578 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_valid_phys = HeedbECGDataset(\n", + " dataset_path=dataset_path,\n", + " split=\"val\",\n", + " sampling_rate=100,\n", + " label_src=\"original_physician\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "2479beb0", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor(100103)" + ] + }, + "execution_count": 26, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_train_phys.labels[:, 0].sum()" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "0c29d791", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "8145485" + ] + }, + "execution_count": 27, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_train_phys.labels.shape[0]" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "ded090ea", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor(0.0123)" + ] + }, + "execution_count": 28, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ds_train_phys.labels[:, 0].sum() / ds_train_phys.labels.shape[0]" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "05a7bf12", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "9070254" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "368191 + 556578 + 8145485" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a7dc6eed", + "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/_eval_probs_bootstrapped.py b/large-user-study/scripts/_eval_probs_bootstrapped.py new file mode 100644 index 0000000..4de0422 --- /dev/null +++ b/large-user-study/scripts/_eval_probs_bootstrapped.py @@ -0,0 +1,142 @@ +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 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) + parser.add_argument("--label-subset", nargs="+") + 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 + 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) + + kwargs = [ + { + "y_test": test_targets[:, src_label_names.index(target_col)], + "y_prob": target_probs[:, src_label_names.index(target_col)], + "n_bootstraps": n_bootstraps, + } + for target_col in 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() + 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/run-protoecgnet-heedb-1D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh new file mode 100755 index 0000000..ece0f1c --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.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-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/gpudata/steven/protoecgnet-heedb +REPO_ROOT=/opt/gpudata/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"}' \ + --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"}' \ + --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"}' \ + --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"}' \ +--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/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..06c7020 --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.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-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/gpudata/steven/protoecgnet-heedb +REPO_ROOT=/opt/gpudata/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"}' \ + --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"}' \ + --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"}' \ + --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"}' \ +--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/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..e875fd9 --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.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-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/gpudata/steven/protoecgnet-heedb +REPO_ROOT=/opt/gpudata/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"}' \ + --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"}' \ + --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"}' \ + --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"}' \ +--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/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..4d09a4a --- /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=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/gpudata/steven/protoecgnet-heedb +REPO_ROOT=/opt/gpudata/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"}' \ + --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/$EXP_NAME/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/$EXP_NAME/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/$EXP_NAME/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"}' \ +--probs-npy $RUN_DIR/$EXP_NAME/train-fusion-classifier/latest/probs.npy \ +--output-path $RUN_DIR/$EXP_NAME + +cp $RUN_DIR/$EXP_NAME/train-fusion-classifier/latest/probs.npy $RUN_DIR/$EXP_NAME/probs.npy diff --git a/protossl/datasets/_heedb_dataset.py b/protossl/datasets/_heedb_dataset.py index 10f4927..fa8ba84 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 @@ -24,6 +26,17 @@ from ._base_dataset import BaseTSDataset, 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 +47,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 +59,17 @@ 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", ): - df = get_heedb_metadata(dataset_path) + df = get_heedb_metadata(dataset_path, heedb_split_type=heedb_split_type) df = df[df["split"] == split].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( @@ -111,7 +128,11 @@ 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=( + split + if heedb_split_type == "by-year" + else f"{split}-{heedb_split_type}" + ), # type: ignore sampling_rate=sampling_rate, ) @@ -120,7 +141,11 @@ 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: @@ -185,11 +210,29 @@ def get_heedb_metadata(heedb_path: str) -> pd.DataFrame: ["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" + 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 how to create splits for HEEDB split type: {heedb_split_type}" + ) full_paths = [] for f, src in zip(df["fpath"], df["source"]): @@ -213,39 +256,41 @@ 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]: + def make_fname_to_code(institution, label_csv, label_col) -> 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") - ) + 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( - df["FileName"], desc=f"Creating {institution} fname to code mapping" + # 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["codes_physician"], + df[label_col], ) } - 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") + _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("mgb", _label_csv, _label_col) + if label_src not in EUH_FNAME_TO_CODE: + EUH_FNAME_TO_CODE[label_src] = make_fname_to_code( + "emory", _label_csv, _label_col + ) 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)} @@ -257,9 +302,9 @@ def make_fname_to_code(institution) -> dict[str, str]: ) ): if institution == "mgb": - codes = MGB_FNAME_TO_CODE.get(fname, "MISSING") + codes = MGB_FNAME_TO_CODE[label_src].get(fname, "MISSING") elif institution == "emory": - codes = EUH_FNAME_TO_CODE.get(fname, "MISSING") + codes = EUH_FNAME_TO_CODE[label_src].get(fname, "MISSING") else: raise ValueError(f"Unknown institution: {institution}") if codes == "MISSING": @@ -280,7 +325,7 @@ def make_fname_to_code(institution) -> dict[str, str]: continue data[meta_idx, label_idx] = 1 print( - f"Of {len(meta)} ECGs and {len(MGB_FNAME_TO_CODE) + len(EUH_FNAME_TO_CODE)} annotations, " + f"Of {len(meta)} ECGs and {len(MGB_FNAME_TO_CODE[label_src]) + len(EUH_FNAME_TO_CODE[label_src])} annotations, " f"{count} matched ({len(meta)-count} ECG annotations were missing and filled with 0s)" ) print("===================================================") diff --git a/protossl/defines.py b/protossl/defines.py index f6ebc64..82f2107 100644 --- a/protossl/defines.py +++ b/protossl/defines.py @@ -11,6 +11,7 @@ "project-prototypes-supervised", "compute-embeddings", "train-classifier", + "train-fusion-classifier", "train-contraster", ] diff --git a/protossl/models/_protoecgnet.py b/protossl/models/_protoecgnet.py new file mode 100644 index 0000000..1f05a86 --- /dev/null +++ b/protossl/models/_protoecgnet.py @@ -0,0 +1,465 @@ +from collections import Counter +from dataclasses import InitVar, dataclass, field +from typing import Literal + +import numpy as np +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 + + +@dataclass +class JointStageLambdas: + lam_clst: float + lam_sep: float + lam_div: float + lam_cntrst: float + + +@dataclass +class ClsStageLambdas: + lam_l1: float + + +LOSS_WEIGHT_T = ( + Literal["protopnet", "protoecgnet"] | JointStageLambdas | ClsStageLambdas +) + + +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, + loss_weights: LOSS_WEIGHT_T = "protoecgnet", + ): + 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) + + if loss_weights == "protopnet": + # original ProtoPNet coefficients + self.lam_clst = 0.8 + self.lam_sep = 0.08 + self.lam_div = 100 + self.lam_cntrst = 0 # no label co-occurrence loss + self.lam_l1 = 1e-4 + elif loss_weights == "protoecgnet": + # ProtoECGNet paper + self.lam_clst = 0.004 + self.lam_sep = 0.0004 + self.lam_div = 250.0 + self.lam_cntrst = 300 + self.lam_l1 = 1e-4 + elif ( + isinstance(loss_weights, JointStageLambdas) + and pipeline_stage == "learn-prototypes-supervised" + ): + self.lam_clst = loss_weights.lam_clst + self.lam_sep = loss_weights.lam_sep + self.lam_div = loss_weights.lam_div + self.lam_cntrst = loss_weights.lam_cntrst + elif ( + isinstance(loss_weights, ClsStageLambdas) + and pipeline_stage == "train-classifier" + ): + self.lam_l1 = loss_weights.lam_l1 + else: + raise ValueError( + f"Unknown how to derive prototype loss weights (got {loss_weights}) for stage (got {pipeline_stage})" + ) + + 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] + + +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, + branches: BRANCHES_T, + loss_weights: dict[str, float] | LOSS_WEIGHT_T = "protoecgnet", + pretrained_weights: str | None = None, + ): + super().__init__() + if pipeline_stage != "train-fusion-classifier": + raise ValueError( + f"Invalid pipeline_stage for ProtoECGNetFusion: {pipeline_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_ppl, 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): + label_to_ppl[label] = branch_ppl + # 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 + assign = torch.zeros(total_labels, total_encoder_dim) # (L, R) + offset = 0 + # the order of the `cls` weights is given by the order of `label_names` + for row, label in enumerate(label_names): + chunk = label_to_ppl[label] + assign[row, offset : offset + chunk] = 1 + offset += chunk + off_class_mask = 1.0 - assign # (L, P') + 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) + + if loss_weights == "protopnet": + # original ProtoPNet coefficients + self.lam_l1 = 1e-4 + elif loss_weights == "protoecgnet": + # ProtoECGNet paper + self.lam_l1 = 1e-4 + elif isinstance(loss_weights, ClsStageLambdas): + self.lam_l1 = loss_weights.lam_l1 + else: + raise ValueError( + f"Unknown how to derive prototype loss weights (got {loss_weights})" + ) + + # 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 forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple[ + dict[str, torch.Tensor], # losses + dict[str, torch.Tensor], # probs + ]: + per_branch_sims = [] + for enc in self.encoders.values(): + per_branch_sims.append(enc(x)) + sims: torch.Tensor = torch.concat(per_branch_sims, dim=1) # (B, R) + sims = sims[:, self.reverse_mask] # type: ignore + 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/protoecgnet_trainer.py b/protossl/protoecgnet_trainer.py new file mode 100644 index 0000000..6af103d --- /dev/null +++ b/protossl/protoecgnet_trainer.py @@ -0,0 +1,688 @@ +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 .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, + LOSS_WEIGHT_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", +] + +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", + } + ), + 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, + loss_weights: LOSS_WEIGHT_T | None = None, + branches: BRANCHES_T | None = None, + ): + super().__init__() + self.lr = None + self.save_hyperparameters() + + 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 + model_kwargs = dict() + if loss_weights is not None: + model_kwargs["loss_weights"] = loss_weights + 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, + **model_kwargs, + ) + 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 + model_kwargs = dict() + if loss_weights is not None: + model_kwargs["loss_weights"] = loss_weights + 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, + **model_kwargs, + ) + 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 + model_kwargs = dict() + if loss_weights is not None: + model_kwargs["loss_weights"] = loss_weights + self.model = ProtoECGNetFusion( + pipeline_stage=pipeline_stage, + label_names=label_names, + label_weights=label_weights, + branches=branches, + pretrained_weights=pretrained_weights, + **model_kwargs, + ) + for param in self.model.encoders.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, + ) + 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 + 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 test_step(self, batch, batch_idx): + loss, _ = self._common_step(batch=batch, stage="test", sync_dist=True) + return loss + + 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 + save_name = "probs.npy" + elif pipeline_stage == "compute-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/ecg-prototype-fm/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": + # 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 = "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, + ) + cli.trainer.predict( + model=cli.model, + datamodule=cli.datamodule, + ) + 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 95151d7..e2016c6 100644 --- a/protossl/trainer.py +++ b/protossl/trainer.py @@ -81,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() 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/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", ) From a15a4e36990f761eb83c552648705013451340ca Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 27 Aug 2026 17:03:21 -0500 Subject: [PATCH 21/33] organize notebooks --- .../confirm-new-heedb-splits.ipynb | 0 .../{ => notebooks}/muse-accuracy.ipynb | 0 large-user-study/notebooks/spoof-heedb.ipynb | 229 ++++++++++++++++++ 3 files changed, 229 insertions(+) rename large-user-study/{ => notebooks}/confirm-new-heedb-splits.ipynb (100%) rename large-user-study/{ => notebooks}/muse-accuracy.ipynb (100%) create mode 100644 large-user-study/notebooks/spoof-heedb.ipynb diff --git a/large-user-study/confirm-new-heedb-splits.ipynb b/large-user-study/notebooks/confirm-new-heedb-splits.ipynb similarity index 100% rename from large-user-study/confirm-new-heedb-splits.ipynb rename to large-user-study/notebooks/confirm-new-heedb-splits.ipynb diff --git a/large-user-study/muse-accuracy.ipynb b/large-user-study/notebooks/muse-accuracy.ipynb similarity index 100% rename from large-user-study/muse-accuracy.ipynb rename to large-user-study/notebooks/muse-accuracy.ipynb 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 +} From fa694e234babf758256bcce7006b8e502bf9be15 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Thu, 27 Aug 2026 17:17:03 -0500 Subject: [PATCH 22/33] paths --- large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh | 4 ++-- large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh | 4 ++-- large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh | 4 ++-- large-user-study/scripts/run-protoecgnet-heedb-fusion.sh | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh index ece0f1c..2e4b9d5 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh @@ -12,8 +12,8 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/protoecgnet-heedb -REPO_ROOT=/opt/gpudata/steven/ProtoSSL/large-user-study +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}" diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh index 06c7020..37438f7 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh @@ -12,8 +12,8 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/protoecgnet-heedb -REPO_ROOT=/opt/gpudata/steven/ProtoSSL/large-user-study +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}" diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh index e875fd9..c0b989e 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh @@ -12,8 +12,8 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/protoecgnet-heedb -REPO_ROOT=/opt/gpudata/steven/ProtoSSL/large-user-study +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}" diff --git a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh index 4d09a4a..752be58 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh @@ -12,8 +12,8 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpudata/steven/protoecgnet-heedb -REPO_ROOT=/opt/gpudata/steven/ProtoSSL/large-user-study +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}" From aad63f3ce0bd943102a4f2b3a601167ac626d083 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Wed, 2 Sep 2026 08:55:50 -0500 Subject: [PATCH 23/33] protoecgnet --- large-user-study/configs/1D-global.yaml | 2 +- large-user-study/configs/2D-global.yaml | 2 +- large-user-study/configs/2D-partial.yaml | 2 +- large-user-study/configs/fusion.yaml | 2 +- .../scripts/_eval_probs_bootstrapped.py | 27 ++++++++++++++++--- .../scripts/run-protoecgnet-heedb-fusion.sh | 6 ++--- protossl/models/_protoecgnet.py | 3 +-- 7 files changed, 31 insertions(+), 13 deletions(-) diff --git a/large-user-study/configs/1D-global.yaml b/large-user-study/configs/1D-global.yaml index c4dfc42..08422bf 100644 --- a/large-user-study/configs/1D-global.yaml +++ b/large-user-study/configs/1D-global.yaml @@ -26,7 +26,7 @@ trainer: - class_path: PredictionWriter max_epochs: 1000 log_every_n_steps: 1 - val_check_interval: 0.25 # affects hyberband pruning + # val_check_interval: 0.25 # set for hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: diff --git a/large-user-study/configs/2D-global.yaml b/large-user-study/configs/2D-global.yaml index d53562b..15f3e09 100644 --- a/large-user-study/configs/2D-global.yaml +++ b/large-user-study/configs/2D-global.yaml @@ -26,7 +26,7 @@ trainer: - class_path: PredictionWriter max_epochs: 1000 log_every_n_steps: 1 - val_check_interval: 0.25 # affects hyberband pruning + # val_check_interval: 0.25 # set for hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: diff --git a/large-user-study/configs/2D-partial.yaml b/large-user-study/configs/2D-partial.yaml index 97f18aa..cb02b2e 100644 --- a/large-user-study/configs/2D-partial.yaml +++ b/large-user-study/configs/2D-partial.yaml @@ -26,7 +26,7 @@ trainer: - class_path: PredictionWriter max_epochs: 1000 log_every_n_steps: 1 - val_check_interval: 0.25 # affects hyberband pruning + # val_check_interval: 0.25 # set for hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: diff --git a/large-user-study/configs/fusion.yaml b/large-user-study/configs/fusion.yaml index 52ed72e..6d99c97 100644 --- a/large-user-study/configs/fusion.yaml +++ b/large-user-study/configs/fusion.yaml @@ -26,7 +26,7 @@ trainer: - class_path: PredictionWriter max_epochs: 1000 log_every_n_steps: 1 - val_check_interval: 0.25 # affects hyberband pruning + # val_check_interval: 0.25 # set for hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: diff --git a/large-user-study/scripts/_eval_probs_bootstrapped.py b/large-user-study/scripts/_eval_probs_bootstrapped.py index 4de0422..0fc64be 100644 --- a/large-user-study/scripts/_eval_probs_bootstrapped.py +++ b/large-user-study/scripts/_eval_probs_bootstrapped.py @@ -6,6 +6,7 @@ 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 @@ -19,7 +20,9 @@ def parse_args(): parser.add_argument("--data-kwargs", required=True) parser.add_argument("--probs-npy", required=True) parser.add_argument("--output-path", required=True) - parser.add_argument("--label-subset", nargs="+") + 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() @@ -96,6 +99,11 @@ def main( 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 @@ -105,14 +113,18 @@ def main( 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 = [ { - "y_test": test_targets[:, src_label_names.index(target_col)], - "y_prob": target_probs[:, src_label_names.index(target_col)], + # 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 target_col in label_names + for i in range(len(label_names)) ] results = pqdm(kwargs, worker, argument_type="kwargs", n_jobs=n_jobs) # type: ignore @@ -131,6 +143,13 @@ def main( 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, diff --git a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh index 752be58..6c324f5 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh @@ -40,15 +40,15 @@ python -m protossl.protoecgnet_trainer \ --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/$EXP_NAME/1D-global/project-prototypes-supervised/latest/proj.ckpt \ + --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/$EXP_NAME/2D-global/project-prototypes-supervised/latest/proj.ckpt \ + --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/$EXP_NAME/2D-partial/project-prototypes-supervised/latest/proj.ckpt \ + --model.branches.pretrained_weights=$RUN_DIR/2D-partial/project-prototypes-supervised/latest/proj.ckpt \ --data.num_workers 8 \ --data.prefetch_factor 4 diff --git a/protossl/models/_protoecgnet.py b/protossl/models/_protoecgnet.py index 1f05a86..66b0db7 100644 --- a/protossl/models/_protoecgnet.py +++ b/protossl/models/_protoecgnet.py @@ -1,8 +1,7 @@ from collections import Counter -from dataclasses import InitVar, dataclass, field +from dataclasses import dataclass, field from typing import Literal -import numpy as np import torch import torch.nn as nn import torch.nn.functional as F From b87eee52caa081bd61edc1af16559745034cb827 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Mon, 14 Sep 2026 11:30:04 -0500 Subject: [PATCH 24/33] branch lambda sweep --- .gitignore | 2 +- large-user-study/configs/1D-global.yaml | 9 +- large-user-study/configs/2D-global.yaml | 9 +- large-user-study/configs/2D-partial.yaml | 9 +- .../configs/lambda-sweep/wandb-sweep.yaml | 56 ++ .../notebooks/muse-accuracy.ipynb | 862 ++++++++---------- protossl/datasets/_heedb_dataset.py | 309 ++++--- protossl/models/_protoecgnet.py | 77 +- protossl/protoecgnet_trainer.py | 95 +- 9 files changed, 735 insertions(+), 693 deletions(-) create mode 100644 large-user-study/configs/lambda-sweep/wandb-sweep.yaml diff --git a/.gitignore b/.gitignore index 59ef6a8..f94b5e5 100644 --- a/.gitignore +++ b/.gitignore @@ -219,4 +219,4 @@ OLD/ slurm-logs/ results/figs/ ecgfounder-checkpoint/ -large-user-study/configs/ppl-sweep/wandb +wandb/ diff --git a/large-user-study/configs/1D-global.yaml b/large-user-study/configs/1D-global.yaml index 08422bf..d0a2850 100644 --- a/large-user-study/configs/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 # set for hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: diff --git a/large-user-study/configs/2D-global.yaml b/large-user-study/configs/2D-global.yaml index 15f3e09..3d288ab 100644 --- a/large-user-study/configs/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 # set for hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: diff --git a/large-user-study/configs/2D-partial.yaml b/large-user-study/configs/2D-partial.yaml index cb02b2e..84ff484 100644 --- a/large-user-study/configs/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 # set for hyberband pruning optimizer: class_path: torch.optim.AdamW init_args: 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/large-user-study/notebooks/muse-accuracy.ipynb b/large-user-study/notebooks/muse-accuracy.ipynb index 97ab010..2da7a57 100644 --- a/large-user-study/notebooks/muse-accuracy.ipynb +++ b/large-user-study/notebooks/muse-accuracy.ipynb @@ -2,24 +2,47 @@ "cells": [ { "cell_type": "code", - "execution_count": 29, + "execution_count": 1, "id": "308c1648", "metadata": {}, "outputs": [], "source": [ "from protossl.datasets import HeedbECGDataset\n", "from protossl.defines import HEEDB_TARGETS\n", - "\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\"" + "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\")" ] }, { "cell_type": "code", "execution_count": 2, + "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 HEEDB_TARGETS]) # 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": 3, "id": "68d44cd8", "metadata": {}, "outputs": [ @@ -28,69 +51,72 @@ "output_type": "stream", "text": [ "================get_heedb_metadata=================\n", - "reading HEEDB metadata\n", - "read MGB data\n", - "read Emory data\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" + "=================make_heedb_labels=================\n", + "reading HEEDB labels from source: /opt/gpudata/ecg/heedb\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ - "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:04<00:00, 2279302.11it/s]\n", - "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 3200781.88it/s]\n", - "Converting code string to labels: 100%|██████████| 368191/368191 [00:03<00:00, 104052.12it/s]\n" + "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:04<00:00, 2293220.64it/s]\n", + "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 2759065.03it/s]\n", + "Converting code string to labels: 100%|██████████| 453512/453512 [00:01<00:00, 287988.11it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Of 368191 ECGs and 11607259 annotations, 368191 matched (0 ECG annotations were missing and filled with 0s)\n", + "saved HEEDB labels to on-disk cache: /home/songs1/.cache/protossl_cache/51ff3555.npy\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" + "=================make_heedb_labels=================\n", + "reading HEEDB labels from source: /opt/gpudata/ecg/heedb\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ - "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:04<00:00, 2147329.70it/s]\n", - "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 3175092.66it/s]\n", - "Converting code string to labels: 100%|██████████| 368191/368191 [00:01<00:00, 242206.99it/s]\n" + "Converting code string to labels: 100%|██████████| 453512/453512 [00:02<00:00, 170936.84it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Of 368191 ECGs and 11607259 annotations, 368191 matched (0 ECG annotations were missing and filled with 0s)\n", + "saved HEEDB labels to on-disk cache: /home/songs1/.cache/protossl_cache/c6af0872.npy\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" + "=================make_heedb_labels=================\n", + "reading HEEDB labels from source: /opt/gpudata/ecg/heedb\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ - "Creating mgb fname to code mapping: 100%|██████████| 10471531/10471531 [00:04<00:00, 2375692.40it/s]\n", - "Creating emory fname to code mapping: 100%|██████████| 974172/974172 [00:00<00:00, 2824799.44it/s]\n", - "Converting code string to labels: 100%|██████████| 368191/368191 [00:05<00:00, 64363.89it/s] \n" + "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:05<00:00, 1896988.96it/s]\n", + "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 3092944.65it/s]\n", + "Converting code string to labels: 100%|██████████| 453512/453512 [00:01<00:00, 298284.00it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Of 368191 ECGs and 11445703 annotations, 368191 matched (0 ECG annotations were missing and filled with 0s)\n", + "saved HEEDB labels to on-disk cache: /home/songs1/.cache/protossl_cache/0a73c965.npy\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", @@ -104,193 +130,95 @@ " split=\"test\",\n", " sampling_rate=100,\n", " label_src=\"original_physician\",\n", + " heedb_split_type=\"by-label\",\n", ")\n", "\n", - "ds_muse = HeedbECGDataset(\n", + "ds_phys_val = HeedbECGDataset(\n", " dataset_path=dataset_path,\n", - " split=\"test\",\n", + " split=\"val\",\n", " sampling_rate=100,\n", - " label_src=\"original_muse\",\n", + " label_src=\"original_physician\",\n", + " heedb_split_type=\"by-label\",\n", ")\n", "\n", - "ds_muse_v24 = HeedbECGDataset(\n", + "ds_muse = HeedbECGDataset(\n", " dataset_path=dataset_path,\n", " split=\"test\",\n", " sampling_rate=100,\n", - " label_src=\"new_muse\",\n", + " label_src=\"original_muse\",\n", + " heedb_split_type=\"by-label\",\n", ")" ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 4, "id": "55a31c9e", "metadata": {}, "outputs": [], "source": [ "y_true = ds_phys.labels.numpy()\n", - "y_pred_muse = ds_muse.labels.numpy()\n", - "y_pred_v24 = ds_muse_v24.labels.numpy()" + "y_true_val = ds_phys_val.labels.numpy()\n", + "y_pred_muse = ds_muse.labels.numpy()" ] }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 5, "id": "a84dca44", "metadata": {}, "outputs": [], "source": [ - "def compute_pred(_y_true, _y_prob):\n", + "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 (_y_prob > optimal_threshold).astype(int)" + " return optimal_threshold" ] }, { "cell_type": "code", - "execution_count": 31, + "execution_count": 6, "id": "be9ee773", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "ANTERIOR INFARCT\n", - "supproto: 0.797\n", - "\n", - "ATRIAL FIBRILLATION\n", - "supproto: 0.962\n", - "\n", - "ATRIAL FLUTTER\n", - "supproto: 0.947\n", - "\n", - "ATRIAL-PACED RHYTHM\n", - "supproto: 0.988\n", - "\n", - "INCOMPLETE RIGHT BUNDLE BRANCH BLOCK\n", - "supproto: 0.967\n", - "\n", - "INFERIOR INFARCT\n", - "supproto: 0.834\n", - "\n", - "LATERAL INFARCT\n", - "supproto: 0.945\n", - "\n", - "LEFT BUNDLE BRANCH BLOCK\n", - "supproto: 0.988\n", - "\n", - "NORMAL SINUS RHYTHM\n", - "supproto: 0.980\n", - "\n", - "PREMATURE ATRIAL COMPLEXES\n", - "supproto: 0.880\n", - "\n", - "PREMATURE VENTRICULAR COMPLEXES\n", - "supproto: 0.875\n", - "\n", - "RIGHT AXIS DEVIATION\n", - "supproto: 0.987\n", - "\n", - "RIGHT BUNDLE BRANCH BLOCK\n", - "supproto: 0.992\n", - "\n", - "SINUS BRADYCARDIA\n", - "supproto: 0.994\n", - "\n", - "SINUS RHYTHM\n", - "supproto: 0.885\n", - "\n", - "SINUS TACHYCARDIA\n", - "supproto: 0.994\n", - "\n", - "VENTRICULAR TACHYCARDIA\n", - "supproto: 0.895\n", - "\n", - "VENTRICULAR-PACED RHYTHM\n", - "supproto: 0.991\n", - "\n", - "WITH 1ST DEGREE AV BLOCK\n", - "supproto: 0.988\n", - "\n", - "WITH SINUS ARRHYTHMIA\n", - "supproto: 0.943\n", - "\n" - ] - } - ], + "outputs": [], "source": [ - "y_prob_supproto = np.load(\"/opt/gpu_working/steven/protossl-ecg-outputs/supproto-heedb/probs.npy\")\n", - "y_prob_protossl = np.load(\"/opt/gpu_working/steven/protossl-ecg-outputs/protossl-heedb/probs.npy\")\n", + "temp = stack_branched_arms()\n", + "y_prob_branched, y_prob_branched_val = temp[\"test\"], temp[\"val\"]\n", + "y_prob_fusion = np.load(run_path / \"fusion/train-fusion-classifier/latest/test_probs.npy\")\n", + "y_prob_fusion_val = np.load(run_path / \"fusion/train-fusion-classifier/latest/val_probs.npy\")\n", "\n", - "y_pred_supproto = np.zeros_like(y_pred_muse)\n", - "y_pred_protossl = np.zeros_like(y_pred_muse)\n", - "results = []\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(HEEDB_TARGETS):\n", - " print(label)\n", - " y_pred_supproto[:, i] = compute_pred(y_true[:, i], y_prob_supproto[:, i])\n", - " y_pred_protossl[:, i] = compute_pred(y_true[:, i], y_prob_protossl[:, i])\n", - " label_results = {\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", - " # \"supproto\": balanced_accuracy_score(y_true[:, i], y_pred_supproto[:, i]),\n", - " # \"protossl\": balanced_accuracy_score(y_true[:, i], y_pred_protossl[:, i]),\n", - " # \"old_muse\": balanced_accuracy_score(y_true[:, i], y_pred_muse[:, i]),\n", - " # \"new_muse\": balanced_accuracy_score(y_true[:, i], y_pred_v24[:, i]),\n", - " \"supproto\": roc_auc_score(y_true[:, i], y_prob_supproto[:, i]),\n", - " # \"protossl\": roc_auc_score(y_true[:, i], y_prob_protossl[:, i]),\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", - " results.append(label_results)\n", - " print(f\"supproto: {label_results['supproto']:0.3f}\")\n", - " # print(f\"protossl: {label_results['protossl']:0.3f}\")\n", - " # print(f\"old_muse: {label_results['old_muse']:0.3f}\")\n", - " # print(f\"new_muse: {label_results['new_muse']:0.3f}\")\n", - " print()" - ] - }, - { - "cell_type": "code", - "execution_count": 32, - "id": "1e01daec", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " label supproto\n", - " ANTERIOR INFARCT 0.797\n", - " ATRIAL FIBRILLATION 0.962\n", - " ATRIAL FLUTTER 0.947\n", - " ATRIAL-PACED RHYTHM 0.988\n", - "INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.967\n", - " INFERIOR INFARCT 0.834\n", - " LATERAL INFARCT 0.945\n", - " LEFT BUNDLE BRANCH BLOCK 0.988\n", - " NORMAL SINUS RHYTHM 0.980\n", - " PREMATURE ATRIAL COMPLEXES 0.880\n", - " PREMATURE VENTRICULAR COMPLEXES 0.875\n", - " RIGHT AXIS DEVIATION 0.987\n", - " RIGHT BUNDLE BRANCH BLOCK 0.992\n", - " SINUS BRADYCARDIA 0.994\n", - " SINUS RHYTHM 0.885\n", - " SINUS TACHYCARDIA 0.994\n", - " VENTRICULAR TACHYCARDIA 0.895\n", - " VENTRICULAR-PACED RHYTHM 0.991\n", - " WITH 1ST DEGREE AV BLOCK 0.988\n", - " WITH SINUS ARRHYTHMIA 0.943\n" - ] - } - ], - "source": [ - "print(pd.DataFrame(results).to_string(float_format=\"%0.3f\", index=False))" + " 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": 33, - "id": "56f18902", + "execution_count": 10, + "id": "23b6c024", "metadata": {}, "outputs": [ { @@ -315,136 +243,164 @@ " \n", " \n", " label\n", - " supproto\n", + " branched: AUROC\n", + " fusion: AUROC\n", " \n", " \n", " \n", " \n", " 0\n", " ANTERIOR INFARCT\n", - " 0.797\n", + " 0.953\n", + " 0.952\n", " \n", " \n", " 1\n", " ATRIAL FIBRILLATION\n", - " 0.962\n", + " 0.958\n", + " 0.958\n", " \n", " \n", " 2\n", " ATRIAL FLUTTER\n", - " 0.947\n", + " 0.952\n", + " 0.951\n", " \n", " \n", " 3\n", " ATRIAL-PACED RHYTHM\n", - " 0.988\n", + " 0.987\n", + " 0.987\n", " \n", " \n", " 4\n", " INCOMPLETE RIGHT BUNDLE BRANCH BLOCK\n", - " 0.967\n", + " 0.962\n", + " 0.962\n", " \n", " \n", " 5\n", " INFERIOR INFARCT\n", - " 0.834\n", + " 0.976\n", + " 0.975\n", " \n", " \n", " 6\n", " LATERAL INFARCT\n", - " 0.945\n", + " 0.957\n", + " 0.956\n", " \n", " \n", " 7\n", " LEFT BUNDLE BRANCH BLOCK\n", - " 0.988\n", + " 0.987\n", + " 0.986\n", " \n", " \n", " 8\n", " NORMAL SINUS RHYTHM\n", - " 0.980\n", + " 0.924\n", + " 0.924\n", " \n", " \n", " 9\n", " PREMATURE ATRIAL COMPLEXES\n", - " 0.880\n", + " 0.901\n", + " 0.908\n", " \n", " \n", " 10\n", " PREMATURE VENTRICULAR COMPLEXES\n", - " 0.875\n", + " 0.893\n", + " 0.893\n", " \n", " \n", " 11\n", " RIGHT AXIS DEVIATION\n", - " 0.987\n", + " 0.983\n", + " 0.982\n", " \n", " \n", " 12\n", " RIGHT BUNDLE BRANCH BLOCK\n", - " 0.992\n", + " 0.985\n", + " 0.984\n", " \n", " \n", " 13\n", " SINUS BRADYCARDIA\n", - " 0.994\n", + " 0.979\n", + " 0.979\n", " \n", " \n", " 14\n", " SINUS RHYTHM\n", - " 0.885\n", + " 0.829\n", + " 0.830\n", " \n", " \n", " 15\n", " SINUS TACHYCARDIA\n", - " 0.994\n", + " 0.980\n", + " 0.980\n", " \n", " \n", " 16\n", " VENTRICULAR TACHYCARDIA\n", - " 0.895\n", + " 0.877\n", + " 0.877\n", " \n", " \n", " 17\n", " VENTRICULAR-PACED RHYTHM\n", - " 0.991\n", + " 0.982\n", + " 0.982\n", " \n", " \n", " 18\n", " WITH 1ST DEGREE AV BLOCK\n", - " 0.988\n", + " 0.977\n", + " 0.977\n", " \n", " \n", " 19\n", " WITH SINUS ARRHYTHMIA\n", - " 0.943\n", + " 0.971\n", + " 0.970\n", + " \n", + " \n", + " 20\n", + " Macro Average\n", + " 0.951\n", + " 0.951\n", " \n", " \n", "\n", "" ], "text/plain": [ - " label supproto\n", - "0 ANTERIOR INFARCT 0.797\n", - "1 ATRIAL FIBRILLATION 0.962\n", - "2 ATRIAL FLUTTER 0.947\n", - "3 ATRIAL-PACED RHYTHM 0.988\n", - "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.967\n", - "5 INFERIOR INFARCT 0.834\n", - "6 LATERAL INFARCT 0.945\n", - "7 LEFT BUNDLE BRANCH BLOCK 0.988\n", - "8 NORMAL SINUS RHYTHM 0.980\n", - "9 PREMATURE ATRIAL COMPLEXES 0.880\n", - "10 PREMATURE VENTRICULAR COMPLEXES 0.875\n", - "11 RIGHT AXIS DEVIATION 0.987\n", - "12 RIGHT BUNDLE BRANCH BLOCK 0.992\n", - "13 SINUS BRADYCARDIA 0.994\n", - "14 SINUS RHYTHM 0.885\n", - "15 SINUS TACHYCARDIA 0.994\n", - "16 VENTRICULAR TACHYCARDIA 0.895\n", - "17 VENTRICULAR-PACED RHYTHM 0.991\n", - "18 WITH 1ST DEGREE AV BLOCK 0.988\n", - "19 WITH SINUS ARRHYTHMIA 0.943" + " label branched: AUROC fusion: AUROC\n", + "0 ANTERIOR INFARCT 0.953 0.952\n", + "1 ATRIAL FIBRILLATION 0.958 0.958\n", + "2 ATRIAL FLUTTER 0.952 0.951\n", + "3 ATRIAL-PACED RHYTHM 0.987 0.987\n", + "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.962 0.962\n", + "5 INFERIOR INFARCT 0.976 0.975\n", + "6 LATERAL INFARCT 0.957 0.956\n", + "7 LEFT BUNDLE BRANCH BLOCK 0.987 0.986\n", + "8 NORMAL SINUS RHYTHM 0.924 0.924\n", + "9 PREMATURE ATRIAL COMPLEXES 0.901 0.908\n", + "10 PREMATURE VENTRICULAR COMPLEXES 0.893 0.893\n", + "11 RIGHT AXIS DEVIATION 0.983 0.982\n", + "12 RIGHT BUNDLE BRANCH BLOCK 0.985 0.984\n", + "13 SINUS BRADYCARDIA 0.979 0.979\n", + "14 SINUS RHYTHM 0.829 0.830\n", + "15 SINUS TACHYCARDIA 0.980 0.980\n", + "16 VENTRICULAR TACHYCARDIA 0.877 0.877\n", + "17 VENTRICULAR-PACED RHYTHM 0.982 0.982\n", + "18 WITH 1ST DEGREE AV BLOCK 0.977 0.977\n", + "19 WITH SINUS ARRHYTHMIA 0.971 0.970\n", + "20 Macro Average 0.951 0.951" ] }, "metadata": {}, @@ -452,286 +408,268 @@ } ], "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(pd.DataFrame(results))" + " display(roc_df)" ] }, { "cell_type": "code", - "execution_count": 34, + "execution_count": 11, "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", + " \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: BalAccbranched: BalAccfusion: BalAcc
0ANTERIOR INFARCT0.9490.8910.891
1ATRIAL FIBRILLATION0.9210.9040.904
2ATRIAL FLUTTER0.8580.8890.888
3ATRIAL-PACED RHYTHM0.8730.9450.946
4INCOMPLETE RIGHT BUNDLE BRANCH BLOCK0.9100.9060.906
5INFERIOR INFARCT0.9720.9190.918
6LATERAL INFARCT0.9730.9080.909
7LEFT BUNDLE BRANCH BLOCK0.9400.9580.957
8NORMAL SINUS RHYTHM0.8750.8570.857
9PREMATURE ATRIAL COMPLEXES0.9100.8350.846
10PREMATURE VENTRICULAR COMPLEXES0.9290.8290.829
11RIGHT AXIS DEVIATION0.9100.9570.956
12RIGHT BUNDLE BRANCH BLOCK0.9690.9620.962
13SINUS BRADYCARDIA0.9570.9560.956
14SINUS RHYTHM0.8510.7550.755
15SINUS TACHYCARDIA0.9580.9540.954
16VENTRICULAR TACHYCARDIA0.6490.7870.789
17VENTRICULAR-PACED RHYTHM0.7380.9460.946
18WITH 1ST DEGREE AV BLOCK0.9670.9370.937
19WITH SINUS ARRHYTHMIA0.9570.9300.929
20Macro Average0.9030.9010.902
\n", + "
" + ], "text/plain": [ - "supproto 0.941648\n", - "dtype: float64" - ] - }, - "execution_count": 34, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "pd.DataFrame(results).set_index(\"label\").mean(axis=0)" - ] - }, - { - "cell_type": "code", - "execution_count": 21, - "id": "d42b88e1", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "tensor(37)" - ] - }, - "execution_count": 21, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ds_phys.labels[:, 0].sum()" - ] - }, - { - "cell_type": "code", - "execution_count": 19, - "id": "7e46f74f", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "tensor(0)" - ] - }, - "execution_count": 19, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ds_muse.labels[:, 0].sum()" - ] - }, - { - "cell_type": "code", - "execution_count": 20, - "id": "fe2e5a2f", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "tensor(21096)" - ] - }, - "execution_count": 20, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ds_muse_v24.labels[:, 0].sum()" - ] - }, - { - "cell_type": "code", - "execution_count": 22, - "id": "a97c7c50", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "torch.Size([368191, 20])" - ] - }, - "execution_count": 22, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ds_phys.labels.shape" - ] - }, - { - "cell_type": "code", - "execution_count": 23, - "id": "4f2bef08", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "=================make_heedb_labels=================\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Converting code string to labels: 100%|██████████| 8145485/8145485 [01:33<00:00, 87111.97it/s] \n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Of 8145485 ECGs and 11607259 annotations, 8145484 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" - ] - } - ], - "source": [ - "ds_train_phys = HeedbECGDataset(\n", - " dataset_path=dataset_path,\n", - " split=\"train\",\n", - " sampling_rate=100,\n", - " label_src=\"original_physician\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 24, - "id": "b1021ed2", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "=================make_heedb_labels=================\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Converting code string to labels: 100%|██████████| 556578/556578 [00:01<00:00, 512100.29it/s]\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Of 556578 ECGs and 11607259 annotations, 556578 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_valid_phys = HeedbECGDataset(\n", - " dataset_path=dataset_path,\n", - " split=\"val\",\n", - " sampling_rate=100,\n", - " label_src=\"original_physician\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 26, - "id": "2479beb0", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "tensor(100103)" - ] - }, - "execution_count": 26, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ds_train_phys.labels[:, 0].sum()" - ] - }, - { - "cell_type": "code", - "execution_count": 27, - "id": "0c29d791", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "8145485" - ] - }, - "execution_count": 27, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ds_train_phys.labels.shape[0]" - ] - }, - { - "cell_type": "code", - "execution_count": 28, - "id": "ded090ea", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "tensor(0.0123)" - ] - }, - "execution_count": 28, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ds_train_phys.labels[:, 0].sum() / ds_train_phys.labels.shape[0]" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "05a7bf12", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "9070254" + " label muse: BalAcc branched: BalAcc \\\n", + "0 ANTERIOR INFARCT 0.949 0.891 \n", + "1 ATRIAL FIBRILLATION 0.921 0.904 \n", + "2 ATRIAL FLUTTER 0.858 0.889 \n", + "3 ATRIAL-PACED RHYTHM 0.873 0.945 \n", + "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.910 0.906 \n", + "5 INFERIOR INFARCT 0.972 0.919 \n", + "6 LATERAL INFARCT 0.973 0.908 \n", + "7 LEFT BUNDLE BRANCH BLOCK 0.940 0.958 \n", + "8 NORMAL SINUS RHYTHM 0.875 0.857 \n", + "9 PREMATURE ATRIAL COMPLEXES 0.910 0.835 \n", + "10 PREMATURE VENTRICULAR COMPLEXES 0.929 0.829 \n", + "11 RIGHT AXIS DEVIATION 0.910 0.957 \n", + "12 RIGHT BUNDLE BRANCH BLOCK 0.969 0.962 \n", + "13 SINUS BRADYCARDIA 0.957 0.956 \n", + "14 SINUS RHYTHM 0.851 0.755 \n", + "15 SINUS TACHYCARDIA 0.958 0.954 \n", + "16 VENTRICULAR TACHYCARDIA 0.649 0.787 \n", + "17 VENTRICULAR-PACED RHYTHM 0.738 0.946 \n", + "18 WITH 1ST DEGREE AV BLOCK 0.967 0.937 \n", + "19 WITH SINUS ARRHYTHMIA 0.957 0.930 \n", + "20 Macro Average 0.903 0.901 \n", + "\n", + " fusion: BalAcc \n", + "0 0.891 \n", + "1 0.904 \n", + "2 0.888 \n", + "3 0.946 \n", + "4 0.906 \n", + "5 0.918 \n", + "6 0.909 \n", + "7 0.957 \n", + "8 0.857 \n", + "9 0.846 \n", + "10 0.829 \n", + "11 0.956 \n", + "12 0.962 \n", + "13 0.956 \n", + "14 0.755 \n", + "15 0.954 \n", + "16 0.789 \n", + "17 0.946 \n", + "18 0.937 \n", + "19 0.929 \n", + "20 0.902 " ] }, - "execution_count": 1, "metadata": {}, - "output_type": "execute_result" + "output_type": "display_data" } ], "source": [ - "368191 + 556578 + 8145485" + "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": "a7dc6eed", + "id": "aa320886", "metadata": {}, "outputs": [], "source": [] diff --git a/protossl/datasets/_heedb_dataset.py b/protossl/datasets/_heedb_dataset.py index fa8ba84..741ddd8 100644 --- a/protossl/datasets/_heedb_dataset.py +++ b/protossl/datasets/_heedb_dataset.py @@ -10,6 +10,7 @@ from tqdm import tqdm from ..defines import ( + CACHE_DIR, HEEDB_EUH_CLIPPED_MEANS, HEEDB_EUH_CLIPPED_STDS, HEEDB_EUH_LOWERS, @@ -62,9 +63,9 @@ def __init__( label_src: LABEL_SRC_T = "original_physician", heedb_split_type: HEEDB_SPLIT_T = "by-year", ): - df = get_heedb_metadata(dataset_path, heedb_split_type=heedb_split_type) + 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) 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( @@ -152,100 +153,110 @@ def get_heedb_metadata( 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"] - ] - - 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!" - ) + # 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") - df.loc[val_test_idxs[:n], "split"] = "val" - df.loc[val_test_idxs[n:], "split"] = "test" - else: - raise ValueError( - f"Unknown how to create splits for HEEDB split type: {heedb_split_type}" + # 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"] + ] - 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 + 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("===================================================") @@ -284,49 +295,87 @@ def make_fname_to_code(institution, label_csv, label_col) -> dict[str, str]: ) } - _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("mgb", _label_csv, _label_col) - if label_src not in EUH_FNAME_TO_CODE: - EUH_FNAME_TO_CODE[label_src] = make_fname_to_code( - "emory", _label_csv, _label_col - ) - 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[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 - 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: + meta_hash = hash_path_list(meta["fpath"]) + identifier = f"HEEDB_labels_{heedb_path.rstrip(os.sep)}_{meta_hash}_{label_subset}_{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 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}") + _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( + "mgb", _label_csv, _label_col + ) + if label_src not in EUH_FNAME_TO_CODE: + EUH_FNAME_TO_CODE[label_src] = make_fname_to_code( + "emory", _label_csv, _label_col + ) + 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[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[label_src]) + len(EUH_FNAME_TO_CODE[label_src])} 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 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/models/_protoecgnet.py b/protossl/models/_protoecgnet.py index 66b0db7..3185059 100644 --- a/protossl/models/_protoecgnet.py +++ b/protossl/models/_protoecgnet.py @@ -1,6 +1,5 @@ from collections import Counter from dataclasses import dataclass, field -from typing import Literal import torch import torch.nn as nn @@ -12,24 +11,6 @@ from .encoders import PrototypeEncoder -@dataclass -class JointStageLambdas: - lam_clst: float - lam_sep: float - lam_div: float - lam_cntrst: float - - -@dataclass -class ClsStageLambdas: - lam_l1: float - - -LOSS_WEIGHT_T = ( - Literal["protopnet", "protoecgnet"] | JointStageLambdas | ClsStageLambdas -) - - class ProtoECGNet(PretrainedMixin, nn.Module): def __init__( self, @@ -45,7 +26,12 @@ def __init__( pretrained_weights: str | None = None, partial_len: int | None = None, partial_overlap: float | None = None, - loss_weights: LOSS_WEIGHT_T = "protoecgnet", + # 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 ( @@ -97,37 +83,11 @@ def __init__( with torch.no_grad(): self.cls.weight.copy_(assign - 0.5 * off_class_mask) - if loss_weights == "protopnet": - # original ProtoPNet coefficients - self.lam_clst = 0.8 - self.lam_sep = 0.08 - self.lam_div = 100 - self.lam_cntrst = 0 # no label co-occurrence loss - self.lam_l1 = 1e-4 - elif loss_weights == "protoecgnet": - # ProtoECGNet paper - self.lam_clst = 0.004 - self.lam_sep = 0.0004 - self.lam_div = 250.0 - self.lam_cntrst = 300 - self.lam_l1 = 1e-4 - elif ( - isinstance(loss_weights, JointStageLambdas) - and pipeline_stage == "learn-prototypes-supervised" - ): - self.lam_clst = loss_weights.lam_clst - self.lam_sep = loss_weights.lam_sep - self.lam_div = loss_weights.lam_div - self.lam_cntrst = loss_weights.lam_cntrst - elif ( - isinstance(loss_weights, ClsStageLambdas) - and pipeline_stage == "train-classifier" - ): - self.lam_l1 = loss_weights.lam_l1 - else: - raise ValueError( - f"Unknown how to derive prototype loss weights (got {loss_weights}) for stage (got {pipeline_stage})" - ) + 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) @@ -290,7 +250,7 @@ def __init__( label_names: list[str], # must have same ordering as input y to forward label_weights: torch.Tensor, branches: BRANCHES_T, - loss_weights: dict[str, float] | LOSS_WEIGHT_T = "protoecgnet", + lam_l1: float = 1e-4, pretrained_weights: str | None = None, ): super().__init__() @@ -409,18 +369,7 @@ def __init__( with torch.no_grad(): self.cls.weight.copy_(assign - 0.5 * off_class_mask) - if loss_weights == "protopnet": - # original ProtoPNet coefficients - self.lam_l1 = 1e-4 - elif loss_weights == "protoecgnet": - # ProtoECGNet paper - self.lam_l1 = 1e-4 - elif isinstance(loss_weights, ClsStageLambdas): - self.lam_l1 = loss_weights.lam_l1 - else: - raise ValueError( - f"Unknown how to derive prototype loss weights (got {loss_weights})" - ) + 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: diff --git a/protossl/protoecgnet_trainer.py b/protossl/protoecgnet_trainer.py index 6af103d..46dbb25 100644 --- a/protossl/protoecgnet_trainer.py +++ b/protossl/protoecgnet_trainer.py @@ -8,17 +8,19 @@ 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, - LOSS_WEIGHT_T, - ProtoECGNet, - ProtoECGNetFusion, -) +from .models._protoecgnet import BRANCHES_T, ProtoECGNet, ProtoECGNetFusion torch.set_float32_matmul_precision("medium") @@ -184,12 +186,35 @@ def __init__( pretrained_weights: str | None = None, partial_len: int | None = None, partial_overlap: float | None = None, - loss_weights: LOSS_WEIGHT_T | 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 @@ -203,9 +228,6 @@ def __init__( 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 - model_kwargs = dict() - if loss_weights is not None: - model_kwargs["loss_weights"] = loss_weights self.model = ProtoECGNet( pipeline_stage=pipeline_stage, backbone_type=backbone_type, @@ -218,7 +240,11 @@ def __init__( pretrained_weights=pretrained_weights, partial_len=partial_len, partial_overlap=partial_overlap, - **model_kwargs, + 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" @@ -248,9 +274,6 @@ def __init__( 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 - model_kwargs = dict() - if loss_weights is not None: - model_kwargs["loss_weights"] = loss_weights self.model = ProtoECGNet( pipeline_stage=pipeline_stage, backbone_type=backbone_type, @@ -262,7 +285,11 @@ def __init__( pretrained_weights=pretrained_weights, partial_len=partial_len, partial_overlap=partial_overlap, - **model_kwargs, + 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 @@ -270,16 +297,13 @@ def __init__( 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 - model_kwargs = dict() - if loss_weights is not None: - model_kwargs["loss_weights"] = loss_weights self.model = ProtoECGNetFusion( pipeline_stage=pipeline_stage, label_names=label_names, label_weights=label_weights, branches=branches, pretrained_weights=pretrained_weights, - **model_kwargs, + lam_l1=lam_l1, ) for param in self.model.encoders.parameters(): param.requires_grad = False @@ -325,6 +349,11 @@ def _common_step( 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) @@ -417,10 +446,20 @@ 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 @@ -526,7 +565,13 @@ def write_on_epoch_end( 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 - save_name = "probs.npy" + 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": assert ( hasattr(pl_module, "prediction_split") @@ -675,9 +720,17 @@ def run(): 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, - datamodule=cli.datamodule, + 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}") From ac5fe9e999f4006d577400513363b92c25597e1b Mon Sep 17 00:00:00 2001 From: StevenSong Date: Wed, 16 Sep 2026 11:30:58 -0500 Subject: [PATCH 25/33] compute fusion embeddings --- ...run-protoecgnet-heedb-1D-global-project.sh | 46 ++++++++++++++++ ...run-protoecgnet-heedb-2D-global-project.sh | 46 ++++++++++++++++ ...un-protoecgnet-heedb-2D-partial-project.sh | 46 ++++++++++++++++ .../run-protoecgnet-heedb-fusion-embed.sh | 53 +++++++++++++++++++ protossl/defines.py | 1 + protossl/lightning_utils.py | 2 +- protossl/models/_protoecgnet.py | 31 ++++++++--- protossl/protoecgnet_trainer.py | 40 +++++++++++++- 8 files changed, 256 insertions(+), 9 deletions(-) create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-2D-global-project.sh create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-2D-partial-project.sh create mode 100755 large-user-study/scripts/run-protoecgnet-heedb-fusion-embed.sh diff --git a/large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh b/large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh new file mode 100755 index 0000000..8731e9c --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=300gb +#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 +BRANCH="1D-global" +EXP_NAME="$BRANCH-tuned" +# best lambda-sweep trial (wandb sweep ibbjt4qe, run 99gpqng0, val_ap=0.5444) +TUNED_TRIAL="v330-prxgttd0" + +# project in the training dataset +python -m protossl.protoecgnet_trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/$BRANCH.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.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/$TUNED_TRIAL/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-global-project.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-global-project.sh new file mode 100755 index 0000000..d8bb3d3 --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-global-project.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=300gb +#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 +BRANCH="2D-global" +EXP_NAME="$BRANCH-tuned" +# best lambda-sweep trial (wandb sweep f5mqrrme, run vuoyocwi, val_ap=0.8759) +TUNED_TRIAL="v337-27ubgqcf" + +# project in the training dataset +python -m protossl.protoecgnet_trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/$BRANCH.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.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/$TUNED_TRIAL/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial-project.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial-project.sh new file mode 100755 index 0000000..ac8019f --- /dev/null +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial-project.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +#SBATCH --cpus-per-task=24 +#SBATCH --mem-per-gpu=300gb +#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 +BRANCH="2D-partial" +EXP_NAME="$BRANCH-tuned" +# best lambda-sweep trial (wandb sweep blqwqt64, run btq9wfqt, val_ap=0.5352) +TUNED_TRIAL="v294-895m7c3p" + +# project in the training dataset +python -m protossl.protoecgnet_trainer \ + --pipeline-stage project-prototypes-supervised \ + --config $REPO_ROOT/configs/$BRANCH.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.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/$TUNED_TRIAL/best.ckpt \ + --data.num_workers 8 \ + --data.prefetch_factor 4 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/protossl/defines.py b/protossl/defines.py index 82f2107..b2e7858 100644 --- a/protossl/defines.py +++ b/protossl/defines.py @@ -12,6 +12,7 @@ "compute-embeddings", "train-classifier", "train-fusion-classifier", + "compute-fusion-embeddings", "train-contraster", ] 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/_protoecgnet.py b/protossl/models/_protoecgnet.py index 3185059..021a68d 100644 --- a/protossl/models/_protoecgnet.py +++ b/protossl/models/_protoecgnet.py @@ -248,16 +248,23 @@ def __init__( *, # enforce kwargs pipeline_stage: STAGE_T, label_names: list[str], # must have same ordering as input y to forward - label_weights: torch.Tensor, + 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 != "train-fusion-classifier": + 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) @@ -375,15 +382,27 @@ def __init__( 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 + 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_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) diff --git a/protossl/protoecgnet_trainer.py b/protossl/protoecgnet_trainer.py index 46dbb25..c2db0d4 100644 --- a/protossl/protoecgnet_trainer.py +++ b/protossl/protoecgnet_trainer.py @@ -30,6 +30,7 @@ "compute-embeddings", "train-classifier", "train-fusion-classifier", + "compute-fusion-embeddings", ] ProtoECGNetTrainerError = ValueError( @@ -137,6 +138,7 @@ def train_dataloader(self): not in { "project-prototypes-supervised", "compute-embeddings", + "compute-fusion-embeddings", } ), pin_memory=True, @@ -307,6 +309,19 @@ def __init__( ) 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}") @@ -412,6 +427,16 @@ def _common_step( # 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( @@ -572,7 +597,10 @@ def write_on_epoch_end( save_name = f"{pl_module.prediction_split}_probs.npy" else: save_name = "probs.npy" - elif pipeline_stage == "compute-embeddings": + 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 @@ -697,7 +725,10 @@ def run(): ) 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": + 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") @@ -705,6 +736,11 @@ def run(): 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( From 7a86512818887ad6138ffc54c3cd6926c1653f0e Mon Sep 17 00:00:00 2001 From: StevenSong Date: Wed, 16 Sep 2026 15:02:30 -0500 Subject: [PATCH 26/33] wip tune fusion classifier --- .../scripts/_tune_fusion_classifier.py | 383 ++++++++++++++++++ ...une-protoecgnet-heedb-fusion-classifier.sh | 45 ++ protossl/models/_protoecgnet.py | 38 +- 3 files changed, 457 insertions(+), 9 deletions(-) create mode 100644 large-user-study/scripts/_tune_fusion_classifier.py create mode 100644 large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh 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..3b9dbde --- /dev/null +++ b/large-user-study/scripts/_tune_fusion_classifier.py @@ -0,0 +1,383 @@ +""" +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. +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 +includes the lam_l1-scaled penalty, so it is not comparable across lam_l1 values +(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 protossl.datasets import ( + HeedbECGDataset, + get_heedb_labels, + get_heedb_metadata, + infer_dataset_class_from_path, +) +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} + + +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("--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) -> 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}") + 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 + ] + return { + "dataset_path": dataset_path, + "data_kwargs": OmegaConf.to_container(data_args.data_kwargs), + "label_subset": label_subset, + "label_names": label_names, + "assign": fusion_prototype_assignment(label_names, branches), # (L, R) + } + + +def load_labels(run_cfg: dict, split: str) -> torch.Tensor: + # same labels as HeedbECGDataset but without instantiating waveforms + data_kwargs = dict(run_cfg["data_kwargs"]) + split_kwargs = {} + if "heedb_split_type" in data_kwargs: + split_kwargs["heedb_split_type"] = data_kwargs.pop("heedb_split_type") + unknown = set(data_kwargs) - {"label_src"} + if len(unknown) > 0: + raise ValueError(f"Unhandled data_kwargs: {unknown}") + df = get_heedb_metadata(run_cfg["dataset_path"], **split_kwargs) + df = df[df["split"] == split].reset_index(drop=True) + labels = get_heedb_labels( + run_cfg["dataset_path"], df, run_cfg["label_subset"], **data_kwargs + ) + return torch.as_tensor(labels) + + +def load_embeds( + embeds_dir: str, + split: str, + device: str, + chunk_rows: int = 500_000, +) -> torch.Tensor: + arr = np.load(os.path.join(embeds_dir, f"{split}_embeds.npy"), mmap_mode="r") + print(f"loading {split} embeds {arr.shape} to {device}") + out = torch.empty(arr.shape, dtype=torch.float32, device=device) + for i in range(0, arr.shape[0], chunk_rows): + out[i : i + chunk_rows].copy_( + torch.from_numpy(np.array(arr[i : i + chunk_rows])) + ) + 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"] + + 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, {} + for epoch in range(max_epochs): + perm = torch.randperm(n_train, device=device, generator=gen) + train_loss_sum = torch.zeros((), device=device) + for start in range(0, n_train, batch_size): + 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) + + Y_train = load_labels(run_cfg, "train") + Y_val = load_labels(run_cfg, "val") + # same as BaseTSDataset.get_label_weights over the train split + per_label_count = Y_train.sum(dim=0) + pos_weight = ((Y_train.shape[0] - per_label_count) / per_label_count).float() + + X_train = load_embeds(args.embeds_dir, "train", device) + X_val = load_embeds(args.embeds_dir, "val", device) + 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]), + } + 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", "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) + 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 + + selection = { + "trial": int(selected["number"]), + "ap_tolerance": args.ap_tolerance, + "best_val_ap": float(best_ap), + "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 main(): + 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) + 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), + ) + 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/tune-protoecgnet-heedb-fusion-classifier.sh b/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh new file mode 100644 index 0000000..eb23675 --- /dev/null +++ b/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh @@ -0,0 +1,45 @@ +#!/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" +EMBEDS_DIR=$RUN_DIR/$EXP_NAME/compute-fusion-embeddings/latest +OUTPUT_DIR=$RUN_DIR/$EXP_NAME/tune-fusion-classifier + +# 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 \ + --n-trials 50 + +python _eval_probs_bootstrapped.py \ +--dataset-path $DATASET_PATH \ +--data-kwargs '{"heedb_split_type": "by-label"}' \ +--probs-npy $OUTPUT_DIR/test_probs.npy \ +--output-path $OUTPUT_DIR + +cp $OUTPUT_DIR/test_probs.npy $RUN_DIR/$EXP_NAME/probs.npy diff --git a/protossl/models/_protoecgnet.py b/protossl/models/_protoecgnet.py index 021a68d..8b06e6b 100644 --- a/protossl/models/_protoecgnet.py +++ b/protossl/models/_protoecgnet.py @@ -242,6 +242,31 @@ def __post_init__(self): 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, @@ -270,7 +295,7 @@ def __init__( total_labels = len(label_names) self.register_buffer("label_weights", label_weights, persistent=False) total_encoder_dim, branch_offset = 0, 0 - encoders, label_to_ppl, label_to_slice = {}, {}, {} + encoders, label_to_slice = {}, {} for branch in branches: if pretrained_weights is None and branch.pretrained_weights is None: raise ValueError( @@ -314,7 +339,6 @@ def __init__( # encoder branch --> cls mapping for branch_idx, label in enumerate(branch.label_subset): - label_to_ppl[label] = branch_ppl # the column slice in the tensor concatenated from all encoders label_to_slice[label] = ( branch_offset + branch_idx * branch_ppl, # start @@ -353,14 +377,10 @@ def __init__( # mask pos/neg prototype class connections # let R be the total number of prototypes across all branches - assign = torch.zeros(total_labels, total_encoder_dim) # (L, R) - offset = 0 # the order of the `cls` weights is given by the order of `label_names` - for row, label in enumerate(label_names): - chunk = label_to_ppl[label] - assign[row, offset : offset + chunk] = 1 - offset += chunk - off_class_mask = 1.0 - assign # (L, P') + 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 From 4090ea500d83353573411cbb61a5bb3f4dfdd209 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sun, 20 Sep 2026 15:37:30 +0000 Subject: [PATCH 27/33] try dropping NSR 2D global --- .gitignore | 1 + .../notebooks/muse-accuracy.ipynb | 368 ++++++------------ .../scripts/_tune_fusion_classifier.py | 192 +++++++-- ...une-protoecgnet-heedb-fusion-classifier.sh | 14 +- 4 files changed, 289 insertions(+), 286 deletions(-) diff --git a/.gitignore b/.gitignore index f94b5e5..002c6b0 100644 --- a/.gitignore +++ b/.gitignore @@ -220,3 +220,4 @@ slurm-logs/ results/figs/ ecgfounder-checkpoint/ wandb/ +.devcontainer/ diff --git a/large-user-study/notebooks/muse-accuracy.ipynb b/large-user-study/notebooks/muse-accuracy.ipynb index 2da7a57..7f2ee89 100644 --- a/large-user-study/notebooks/muse-accuracy.ipynb +++ b/large-user-study/notebooks/muse-accuracy.ipynb @@ -18,19 +18,27 @@ "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\")" + "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": 2, + "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 HEEDB_TARGETS]) # indices to normalize to fusion order\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", @@ -42,7 +50,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 2, "id": "68d44cd8", "metadata": {}, "outputs": [ @@ -55,67 +63,24 @@ "reading HEEDB metadata from on-disk cache: /home/songs1/.cache/protossl_cache/abac20ad.csv\n", "===================================================\n", "=================make_heedb_labels=================\n", - "reading HEEDB labels from source: /opt/gpudata/ecg/heedb\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:04<00:00, 2293220.64it/s]\n", - "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 2759065.03it/s]\n", - "Converting code string to labels: 100%|██████████| 453512/453512 [00:01<00:00, 287988.11it/s]\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "saved HEEDB labels to on-disk cache: /home/songs1/.cache/protossl_cache/51ff3555.npy\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 source: /opt/gpudata/ecg/heedb\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Converting code string to labels: 100%|██████████| 453512/453512 [00:02<00:00, 170936.84it/s]\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "saved HEEDB labels to on-disk cache: /home/songs1/.cache/protossl_cache/c6af0872.npy\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 source: /opt/gpudata/ecg/heedb\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Creating mgb fname to code mapping: 100%|██████████| 10608415/10608415 [00:05<00:00, 1896988.96it/s]\n", - "Creating emory fname to code mapping: 100%|██████████| 998844/998844 [00:00<00:00, 3092944.65it/s]\n", - "Converting code string to labels: 100%|██████████| 453512/453512 [00:01<00:00, 298284.00it/s]\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "saved HEEDB labels to on-disk cache: /home/songs1/.cache/protossl_cache/0a73c965.npy\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", @@ -152,19 +117,19 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 3, "id": "55a31c9e", "metadata": {}, "outputs": [], "source": [ - "y_true = ds_phys.labels.numpy()\n", - "y_true_val = ds_phys_val.labels.numpy()\n", - "y_pred_muse = ds_muse.labels.numpy()" + "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": 5, + "execution_count": 4, "id": "a84dca44", "metadata": {}, "outputs": [], @@ -180,35 +145,36 @@ }, { "cell_type": "code", - "execution_count": 6, + "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(run_path / \"fusion/train-fusion-classifier/latest/test_probs.npy\")\n", - "y_prob_fusion_val = np.load(run_path / \"fusion/train-fusion-classifier/latest/val_probs.npy\")\n", + "# 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_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(HEEDB_TARGETS):\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", + "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", + " # \"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", + " # \"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", @@ -217,7 +183,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 6, "id": "23b6c024", "metadata": {}, "outputs": [ @@ -243,7 +209,6 @@ " \n", " \n", " label\n", - " branched: AUROC\n", " fusion: AUROC\n", " \n", " \n", @@ -252,155 +217,128 @@ " 0\n", " ANTERIOR INFARCT\n", " 0.953\n", - " 0.952\n", " \n", " \n", " 1\n", " ATRIAL FIBRILLATION\n", - " 0.958\n", - " 0.958\n", + " 0.959\n", " \n", " \n", " 2\n", " ATRIAL FLUTTER\n", - " 0.952\n", - " 0.951\n", + " 0.953\n", " \n", " \n", " 3\n", " ATRIAL-PACED RHYTHM\n", - " 0.987\n", - " 0.987\n", + " 0.986\n", " \n", " \n", " 4\n", " INCOMPLETE RIGHT BUNDLE BRANCH BLOCK\n", " 0.962\n", - " 0.962\n", " \n", " \n", " 5\n", " INFERIOR INFARCT\n", " 0.976\n", - " 0.975\n", " \n", " \n", " 6\n", " LATERAL INFARCT\n", " 0.957\n", - " 0.956\n", " \n", " \n", " 7\n", " LEFT BUNDLE BRANCH BLOCK\n", " 0.987\n", - " 0.986\n", " \n", " \n", " 8\n", - " NORMAL SINUS RHYTHM\n", - " 0.924\n", - " 0.924\n", - " \n", - " \n", - " 9\n", " PREMATURE ATRIAL COMPLEXES\n", - " 0.901\n", - " 0.908\n", + " 0.912\n", " \n", " \n", - " 10\n", + " 9\n", " PREMATURE VENTRICULAR COMPLEXES\n", - " 0.893\n", - " 0.893\n", + " 0.896\n", " \n", " \n", - " 11\n", + " 10\n", " RIGHT AXIS DEVIATION\n", - " 0.983\n", " 0.982\n", " \n", " \n", - " 12\n", + " 11\n", " RIGHT BUNDLE BRANCH BLOCK\n", - " 0.985\n", - " 0.984\n", + " 0.986\n", " \n", " \n", - " 13\n", + " 12\n", " SINUS BRADYCARDIA\n", - " 0.979\n", - " 0.979\n", + " 0.981\n", " \n", " \n", - " 14\n", + " 13\n", " SINUS RHYTHM\n", - " 0.829\n", - " 0.830\n", + " 0.833\n", " \n", " \n", - " 15\n", + " 14\n", " SINUS TACHYCARDIA\n", - " 0.980\n", - " 0.980\n", + " 0.981\n", " \n", " \n", - " 16\n", + " 15\n", " VENTRICULAR TACHYCARDIA\n", - " 0.877\n", - " 0.877\n", + " 0.856\n", " \n", " \n", - " 17\n", + " 16\n", " VENTRICULAR-PACED RHYTHM\n", " 0.982\n", - " 0.982\n", " \n", " \n", - " 18\n", + " 17\n", " WITH 1ST DEGREE AV BLOCK\n", - " 0.977\n", - " 0.977\n", + " 0.978\n", " \n", " \n", - " 19\n", + " 18\n", " WITH SINUS ARRHYTHMIA\n", - " 0.971\n", - " 0.970\n", + " 0.973\n", " \n", " \n", - " 20\n", + " 19\n", " Macro Average\n", - " 0.951\n", - " 0.951\n", + " 0.952\n", " \n", " \n", "\n", "" ], "text/plain": [ - " label branched: AUROC fusion: AUROC\n", - "0 ANTERIOR INFARCT 0.953 0.952\n", - "1 ATRIAL FIBRILLATION 0.958 0.958\n", - "2 ATRIAL FLUTTER 0.952 0.951\n", - "3 ATRIAL-PACED RHYTHM 0.987 0.987\n", - "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.962 0.962\n", - "5 INFERIOR INFARCT 0.976 0.975\n", - "6 LATERAL INFARCT 0.957 0.956\n", - "7 LEFT BUNDLE BRANCH BLOCK 0.987 0.986\n", - "8 NORMAL SINUS RHYTHM 0.924 0.924\n", - "9 PREMATURE ATRIAL COMPLEXES 0.901 0.908\n", - "10 PREMATURE VENTRICULAR COMPLEXES 0.893 0.893\n", - "11 RIGHT AXIS DEVIATION 0.983 0.982\n", - "12 RIGHT BUNDLE BRANCH BLOCK 0.985 0.984\n", - "13 SINUS BRADYCARDIA 0.979 0.979\n", - "14 SINUS RHYTHM 0.829 0.830\n", - "15 SINUS TACHYCARDIA 0.980 0.980\n", - "16 VENTRICULAR TACHYCARDIA 0.877 0.877\n", - "17 VENTRICULAR-PACED RHYTHM 0.982 0.982\n", - "18 WITH 1ST DEGREE AV BLOCK 0.977 0.977\n", - "19 WITH SINUS ARRHYTHMIA 0.971 0.970\n", - "20 Macro Average 0.951 0.951" + " 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": {}, @@ -420,7 +358,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 7, "id": "c2c33568", "metadata": {}, "outputs": [ @@ -447,7 +385,6 @@ " \n", " label\n", " muse: BalAcc\n", - " branched: BalAcc\n", " fusion: BalAcc\n", " \n", " \n", @@ -456,49 +393,42 @@ " 0\n", " ANTERIOR INFARCT\n", " 0.949\n", - " 0.891\n", - " 0.891\n", + " 0.893\n", " \n", " \n", " 1\n", " ATRIAL FIBRILLATION\n", " 0.921\n", - " 0.904\n", - " 0.904\n", + " 0.905\n", " \n", " \n", " 2\n", " ATRIAL FLUTTER\n", " 0.858\n", - " 0.889\n", " 0.888\n", " \n", " \n", " 3\n", " ATRIAL-PACED RHYTHM\n", " 0.873\n", - " 0.945\n", - " 0.946\n", + " 0.943\n", " \n", " \n", " 4\n", " INCOMPLETE RIGHT BUNDLE BRANCH BLOCK\n", " 0.910\n", - " 0.906\n", - " 0.906\n", + " 0.908\n", " \n", " \n", " 5\n", " INFERIOR INFARCT\n", " 0.972\n", - " 0.919\n", " 0.918\n", " \n", " \n", " 6\n", " LATERAL INFARCT\n", " 0.973\n", - " 0.908\n", " 0.909\n", " \n", " \n", @@ -506,149 +436,105 @@ " LEFT BUNDLE BRANCH BLOCK\n", " 0.940\n", " 0.958\n", - " 0.957\n", " \n", " \n", " 8\n", - " NORMAL SINUS RHYTHM\n", - " 0.875\n", - " 0.857\n", - " 0.857\n", - " \n", - " \n", - " 9\n", " PREMATURE ATRIAL COMPLEXES\n", " 0.910\n", - " 0.835\n", - " 0.846\n", + " 0.850\n", " \n", " \n", - " 10\n", + " 9\n", " PREMATURE VENTRICULAR COMPLEXES\n", " 0.929\n", - " 0.829\n", - " 0.829\n", + " 0.831\n", " \n", " \n", - " 11\n", + " 10\n", " RIGHT AXIS DEVIATION\n", " 0.910\n", - " 0.957\n", " 0.956\n", " \n", " \n", - " 12\n", + " 11\n", " RIGHT BUNDLE BRANCH BLOCK\n", " 0.969\n", " 0.962\n", - " 0.962\n", " \n", " \n", - " 13\n", + " 12\n", " SINUS BRADYCARDIA\n", " 0.957\n", " 0.956\n", - " 0.956\n", " \n", " \n", - " 14\n", + " 13\n", " SINUS RHYTHM\n", " 0.851\n", - " 0.755\n", - " 0.755\n", + " 0.757\n", " \n", " \n", - " 15\n", + " 14\n", " SINUS TACHYCARDIA\n", " 0.958\n", " 0.954\n", - " 0.954\n", " \n", " \n", - " 16\n", + " 15\n", " VENTRICULAR TACHYCARDIA\n", " 0.649\n", - " 0.787\n", - " 0.789\n", + " 0.773\n", " \n", " \n", - " 17\n", + " 16\n", " VENTRICULAR-PACED RHYTHM\n", " 0.738\n", - " 0.946\n", - " 0.946\n", + " 0.948\n", " \n", " \n", - " 18\n", + " 17\n", " WITH 1ST DEGREE AV BLOCK\n", " 0.967\n", - " 0.937\n", - " 0.937\n", + " 0.938\n", " \n", " \n", - " 19\n", + " 18\n", " WITH SINUS ARRHYTHMIA\n", " 0.957\n", - " 0.930\n", - " 0.929\n", + " 0.931\n", " \n", " \n", - " 20\n", + " 19\n", " Macro Average\n", - " 0.903\n", - " 0.901\n", - " 0.902\n", + " 0.905\n", + " 0.904\n", " \n", " \n", "\n", "" ], "text/plain": [ - " label muse: BalAcc branched: BalAcc \\\n", - "0 ANTERIOR INFARCT 0.949 0.891 \n", - "1 ATRIAL FIBRILLATION 0.921 0.904 \n", - "2 ATRIAL FLUTTER 0.858 0.889 \n", - "3 ATRIAL-PACED RHYTHM 0.873 0.945 \n", - "4 INCOMPLETE RIGHT BUNDLE BRANCH BLOCK 0.910 0.906 \n", - "5 INFERIOR INFARCT 0.972 0.919 \n", - "6 LATERAL INFARCT 0.973 0.908 \n", - "7 LEFT BUNDLE BRANCH BLOCK 0.940 0.958 \n", - "8 NORMAL SINUS RHYTHM 0.875 0.857 \n", - "9 PREMATURE ATRIAL COMPLEXES 0.910 0.835 \n", - "10 PREMATURE VENTRICULAR COMPLEXES 0.929 0.829 \n", - "11 RIGHT AXIS DEVIATION 0.910 0.957 \n", - "12 RIGHT BUNDLE BRANCH BLOCK 0.969 0.962 \n", - "13 SINUS BRADYCARDIA 0.957 0.956 \n", - "14 SINUS RHYTHM 0.851 0.755 \n", - "15 SINUS TACHYCARDIA 0.958 0.954 \n", - "16 VENTRICULAR TACHYCARDIA 0.649 0.787 \n", - "17 VENTRICULAR-PACED RHYTHM 0.738 0.946 \n", - "18 WITH 1ST DEGREE AV BLOCK 0.967 0.937 \n", - "19 WITH SINUS ARRHYTHMIA 0.957 0.930 \n", - "20 Macro Average 0.903 0.901 \n", - "\n", - " fusion: BalAcc \n", - "0 0.891 \n", - "1 0.904 \n", - "2 0.888 \n", - "3 0.946 \n", - "4 0.906 \n", - "5 0.918 \n", - "6 0.909 \n", - "7 0.957 \n", - "8 0.857 \n", - "9 0.846 \n", - "10 0.829 \n", - "11 0.956 \n", - "12 0.962 \n", - "13 0.956 \n", - "14 0.755 \n", - "15 0.954 \n", - "16 0.789 \n", - "17 0.946 \n", - "18 0.937 \n", - "19 0.929 \n", - "20 0.902 " + " 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": {}, diff --git a/large-user-study/scripts/_tune_fusion_classifier.py b/large-user-study/scripts/_tune_fusion_classifier.py index 3b9dbde..434d4cc 100644 --- a/large-user-study/scripts/_tune_fusion_classifier.py +++ b/large-user-study/scripts/_tune_fusion_classifier.py @@ -7,13 +7,21 @@ 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. +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 -includes the lam_l1-scaled penalty, so it is not comparable across lam_l1 values -(it is still used for the LR schedule within a trial, as in the Lightning config). +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 @@ -29,19 +37,22 @@ multilabel_auroc, multilabel_average_precision, ) +from tqdm import trange -from protossl.datasets import ( - HeedbECGDataset, - get_heedb_labels, - get_heedb_metadata, - infer_dataset_class_from_path, -) +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} +BASELINE_PARAMS = { + "lam_l1": 1e-4, + "lr": 1e-3, + "weight_decay": 0.01, + "batch_size": 512, + "pos_weight_exp": 1.0, +} def parse_args(): @@ -52,6 +63,18 @@ def parse_args(): 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) @@ -75,7 +98,7 @@ def parse_args(): return parser.parse_args() -def load_run_config(embeds_dir: str) -> dict: +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 @@ -83,6 +106,7 @@ def load_run_config(embeds_dir: str) -> dict: 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) @@ -91,45 +115,93 @@ def load_run_config(embeds_dir: str) -> dict: 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_labels(run_cfg: dict, split: str) -> torch.Tensor: - # same labels as HeedbECGDataset but without instantiating waveforms - data_kwargs = dict(run_cfg["data_kwargs"]) - split_kwargs = {} - if "heedb_split_type" in data_kwargs: - split_kwargs["heedb_split_type"] = data_kwargs.pop("heedb_split_type") - unknown = set(data_kwargs) - {"label_src"} - if len(unknown) > 0: - raise ValueError(f"Unhandled data_kwargs: {unknown}") - df = get_heedb_metadata(run_cfg["dataset_path"], **split_kwargs) - df = df[df["split"] == split].reset_index(drop=True) - labels = get_heedb_labels( - run_cfg["dataset_path"], df, run_cfg["label_subset"], **data_kwargs +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"], ) - return torch.as_tensor(labels) 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") - print(f"loading {split} embeds {arr.shape} to {device}") - out = torch.empty(arr.shape, dtype=torch.float32, device=device) + 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): - out[i : i + chunk_rows].copy_( - torch.from_numpy(np.array(arr[i : i + 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 @@ -174,6 +246,9 @@ def train_head( 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) @@ -185,10 +260,11 @@ def train_head( ) best_ap, best_epoch, best_W, best_metrics = -1.0, -1, None, {} - for epoch in range(max_epochs): + 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 range(0, n_train, batch_size): + 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( @@ -252,14 +328,13 @@ def tune(args, run_cfg: dict, study: optuna.Study, trials_dir: str): device = args.device assign = run_cfg["assign"].to(device) - Y_train = load_labels(run_cfg, "train") - Y_val = load_labels(run_cfg, "val") - # same as BaseTSDataset.get_label_weights over the train split - per_label_count = Y_train.sum(dim=0) - pos_weight = ((Y_train.shape[0] - per_label_count) / per_label_count).float() + 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) - X_val = load_embeds(args.embeds_dir, "val", device) + 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}" @@ -272,6 +347,7 @@ def objective(trial: optuna.Trial) -> float: "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( @@ -322,7 +398,8 @@ def export(args, run_cfg: dict, study: optuna.Study, trials_dir: str): selected = candidates.sort_values("off_class_frac").iloc[0] summary_cols = [ "number", "value", "val_auroc", "lam_l1", "lr", "weight_decay", - "batch_size", "off_class_frac", "off_class_near_zero_frac", "best_epoch", + "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( @@ -335,17 +412,24 @@ def export(args, run_cfg: dict, study: optuna.Study, trials_dir: str): 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) + 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) @@ -358,13 +442,34 @@ def export(args, run_cfg: dict, study: optuna.Study, trials_dir: 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) + 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')}", @@ -373,6 +478,7 @@ def main(): 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 diff --git a/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh b/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh index eb23675..0aa061e 100644 --- a/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh +++ b/large-user-study/scripts/tune-protoecgnet-heedb-fusion-classifier.sh @@ -26,20 +26,30 @@ 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 +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 -cp $OUTPUT_DIR/test_probs.npy $RUN_DIR/$EXP_NAME/probs.npy +# 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 From f2efcc12ce72098c4ac3d362dc82a920d38d7f60 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Wed, 23 Sep 2026 18:53:40 +0000 Subject: [PATCH 28/33] cache heedb labels using target set --- protossl/datasets/_heedb_dataset.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/protossl/datasets/_heedb_dataset.py b/protossl/datasets/_heedb_dataset.py index 741ddd8..7e1301e 100644 --- a/protossl/datasets/_heedb_dataset.py +++ b/protossl/datasets/_heedb_dataset.py @@ -296,7 +296,8 @@ def make_fname_to_code(institution, label_csv, label_col) -> dict[str, str]: } meta_hash = hash_path_list(meta["fpath"]) - identifier = f"HEEDB_labels_{heedb_path.rstrip(os.sep)}_{meta_hash}_{label_subset}_{label_src}" + 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") From 1a2e80f6a7bb55dc873fb888348fb86c9d0b33ef Mon Sep 17 00:00:00 2001 From: StevenSong Date: Wed, 23 Sep 2026 17:36:27 -0500 Subject: [PATCH 29/33] run overread corrected heedb, run audio downstream --- large-user-study/configs/2D-global.yaml | 2 +- .../scripts/_cache_heedb_overread.py | 18 ++ .../queue-protoecgnet-heedb-overread.sh | 18 ++ .../scripts/run-cache-heedb-overread.sh | 20 ++ .../run-protoecgnet-heedb-1D-global.sh | 12 +- .../run-protoecgnet-heedb-2D-global.sh | 12 +- .../run-protoecgnet-heedb-2D-partial.sh | 12 +- .../scripts/run-protoecgnet-heedb-fusion.sh | 8 +- protossl/datasets/_base_dataset.py | 20 +- protossl/datasets/_heedb_dataset.py | 183 ++++++++++++++---- scripts/audio/3-run-protossl-audioset-pila.sh | 4 +- .../audio/4-run-labsup-proto-audioset-rila.sh | 4 +- scripts/audio/queue-experiments.sh | 2 +- .../voxceleb-2-run-labsup-proto-direct.sh | 2 +- .../voxceleb-3-run-protossl-audioset-pila.sh | 4 +- ...xceleb-4-run-labsup-proto-audioset-rila.sh | 4 +- 16 files changed, 252 insertions(+), 73 deletions(-) create mode 100644 large-user-study/scripts/_cache_heedb_overread.py create mode 100755 large-user-study/scripts/queue-protoecgnet-heedb-overread.sh create mode 100755 large-user-study/scripts/run-cache-heedb-overread.sh diff --git a/large-user-study/configs/2D-global.yaml b/large-user-study/configs/2D-global.yaml index 3d288ab..0a31a36 100644 --- a/large-user-study/configs/2D-global.yaml +++ b/large-user-study/configs/2D-global.yaml @@ -42,7 +42,7 @@ model: backbone_type: resnet18 conv_type: 2D prototype_type: global - n_prototypes_per_label: 256 + n_prototypes_per_label: 50 # pretrained_weights: /path/to/weights data: class_path: LitData 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/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 index 2e4b9d5..9786ba5 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=500gb +#SBATCH --mem-per-gpu=450gb #SBATCH --gpus-per-node=1 #SBATCH --nodes=1 #SBATCH -w kg35-nvl01 @@ -12,7 +12,7 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpu_working/steven/protoecgnet-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 @@ -37,7 +37,7 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ --data.num_workers 8 \ --data.prefetch_factor 4 @@ -48,7 +48,7 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --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 @@ -59,14 +59,14 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --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"}' \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ --probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ --output-path $RUN_DIR/$EXP_NAME diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh index 37438f7..6286a53 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=500gb +#SBATCH --mem-per-gpu=450gb #SBATCH --gpus-per-node=1 #SBATCH --nodes=1 #SBATCH -w kg35-nvl01 @@ -12,7 +12,7 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpu_working/steven/protoecgnet-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 @@ -37,7 +37,7 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ --data.num_workers 8 \ --data.prefetch_factor 4 @@ -48,7 +48,7 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --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 @@ -59,14 +59,14 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --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"}' \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ --probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ --output-path $RUN_DIR/$EXP_NAME diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh index c0b989e..bac4280 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=500gb +#SBATCH --mem-per-gpu=450gb #SBATCH --gpus-per-node=1 #SBATCH --nodes=1 #SBATCH -w kg35-nvl01 @@ -12,7 +12,7 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpu_working/steven/protoecgnet-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 @@ -37,7 +37,7 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --data.data_kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ --data.num_workers 8 \ --data.prefetch_factor 4 @@ -48,7 +48,7 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --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 @@ -59,14 +59,14 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --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"}' \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ --probs-npy $RUN_DIR/$EXP_NAME/train-classifier/latest/probs.npy \ --output-path $RUN_DIR/$EXP_NAME diff --git a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh index 6c324f5..1812fcb 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=500gb +#SBATCH --mem-per-gpu=450gb #SBATCH --gpus-per-node=1 #SBATCH --nodes=1 #SBATCH -w kg35-nvl01 @@ -12,7 +12,7 @@ set -e DATASET_PATH=/opt/gpudata/ecg/heedb -RUN_DIR=/opt/gpu_working/steven/protoecgnet-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 @@ -36,7 +36,7 @@ python -m protossl.protoecgnet_trainer \ --trainer.logger.save_dir $RUN_DIR \ --trainer.logger.name $EXP_NAME \ --data.dataset_path $DATASET_PATH \ - --data.data_kwargs '{"heedb_split_type": "by-label"}' \ + --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 \ @@ -54,7 +54,7 @@ python -m protossl.protoecgnet_trainer \ python _eval_probs_bootstrapped.py \ --dataset-path $DATASET_PATH \ ---data-kwargs '{"heedb_split_type": "by-label"}' \ +--data-kwargs '{"heedb_split_type": "by-label", "drop_blank_overread": True}' \ --probs-npy $RUN_DIR/$EXP_NAME/train-fusion-classifier/latest/probs.npy \ --output-path $RUN_DIR/$EXP_NAME 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 7e1301e..68c7e62 100644 --- a/protossl/datasets/_heedb_dataset.py +++ b/protossl/datasets/_heedb_dataset.py @@ -24,7 +24,12 @@ 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 @@ -62,10 +67,25 @@ def __init__( 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, ): full_df = get_heedb_metadata(dataset_path, heedb_split_type=heedb_split_type) 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( @@ -107,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") @@ -129,11 +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 - if heedb_split_type == "by-year" - else f"{split}-{heedb_split_type}" - ), # type: ignore + split=cache_split, # type: ignore sampling_rate=sampling_rate, ) @@ -275,26 +319,6 @@ def get_heedb_labels( validate_label_subset(label_subset, list(HEEDB_TARGETS)) targets = {label: HEEDB_TARGETS[label] for label in label_subset} - def make_fname_to_code(institution, label_csv, label_col) -> 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], - ) - } - 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}" @@ -310,15 +334,7 @@ def make_fname_to_code(institution, label_csv, label_col) -> dict[str, str]: print("replaying below stats from cache:") else: print(f"reading HEEDB labels from source: {heedb_path}") - _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( - "mgb", _label_csv, _label_col - ) - if label_src not in EUH_FNAME_TO_CODE: - EUH_FNAME_TO_CODE[label_src] = make_fname_to_code( - "emory", _label_csv, _label_col - ) + 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)} @@ -374,6 +390,103 @@ def make_fname_to_code(institution, label_csv, label_col) -> dict[str, str]: 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): diff --git a/scripts/audio/3-run-protossl-audioset-pila.sh b/scripts/audio/3-run-protossl-audioset-pila.sh index c6f72b8..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,7 +28,7 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="protossl-audioset-pila" -PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2.BAK") +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2") python -m protossl.trainer \ --seed_everything $SEED \ diff --git a/scripts/audio/4-run-labsup-proto-audioset-rila.sh b/scripts/audio/4-run-labsup-proto-audioset-rila.sh index 4ab7b29..a2e578a 100755 --- a/scripts/audio/4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/4-run-labsup-proto-audioset-rila.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,7 +28,7 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="labsup-proto-audioset-rila" -PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2.BAK") +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2") python -m protossl.trainer \ --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ diff --git a/scripts/audio/queue-experiments.sh b/scripts/audio/queue-experiments.sh index f401a72..6b746c1 100644 --- a/scripts/audio/queue-experiments.sh +++ b/scripts/audio/queue-experiments.sh @@ -9,7 +9,7 @@ scripts=( voxceleb_scripts=( # "1-run-blackbox-direct.sh" - # "voxceleb-2-run-labsup-proto-direct.sh" + "voxceleb-2-run-labsup-proto-direct.sh" "voxceleb-3-run-protossl-audioset-pila.sh" "voxceleb-4-run-labsup-proto-audioset-rila.sh" ) diff --git a/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh b/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh index 719d1fe..4f70c6d 100755 --- a/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh +++ b/scripts/audio/voxceleb-2-run-labsup-proto-direct.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 diff --git a/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh b/scripts/audio/voxceleb-3-run-protossl-audioset-pila.sh index de0e0f8..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 @@ -32,7 +32,7 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="protossl-audioset-pila" -PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2.BAK") +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../protossl-audioset-v2") python -m protossl.trainer \ --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ diff --git a/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh b/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh index 036d455..d3b0d92 100755 --- a/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.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 @@ -32,7 +32,7 @@ cd $REPO_ROOT/scripts/audio # experiment parameters EXP_NAME="labsup-proto-audioset-rila" -PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2.BAK") +PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2") python -m protossl.trainer \ --config $REPO_ROOT/configs/audio/target-guided-$PPL.yaml \ From cf2180a41f4c47a7f1ffd6004096cc9dbe13e1d9 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sun, 27 Sep 2026 11:20:01 -0500 Subject: [PATCH 30/33] tweaks --- ...run-protoecgnet-heedb-1D-global-project.sh | 46 ------------------- .../run-protoecgnet-heedb-1D-global.sh | 5 +- ...run-protoecgnet-heedb-2D-global-project.sh | 46 ------------------- .../run-protoecgnet-heedb-2D-global.sh | 5 +- ...un-protoecgnet-heedb-2D-partial-project.sh | 46 ------------------- .../run-protoecgnet-heedb-2D-partial.sh | 5 +- .../scripts/run-protoecgnet-heedb-fusion.sh | 4 +- 7 files changed, 11 insertions(+), 146 deletions(-) delete mode 100755 large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh delete mode 100755 large-user-study/scripts/run-protoecgnet-heedb-2D-global-project.sh delete mode 100755 large-user-study/scripts/run-protoecgnet-heedb-2D-partial-project.sh diff --git a/large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh b/large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh deleted file mode 100755 index 8731e9c..0000000 --- a/large-user-study/scripts/run-protoecgnet-heedb-1D-global-project.sh +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/bash - -#SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=300gb -#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 -BRANCH="1D-global" -EXP_NAME="$BRANCH-tuned" -# best lambda-sweep trial (wandb sweep ibbjt4qe, run 99gpqng0, val_ap=0.5444) -TUNED_TRIAL="v330-prxgttd0" - -# project in the training dataset -python -m protossl.protoecgnet_trainer \ - --pipeline-stage project-prototypes-supervised \ - --config $REPO_ROOT/configs/$BRANCH.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.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/$TUNED_TRIAL/best.ckpt \ - --data.num_workers 8 \ - --data.prefetch_factor 4 diff --git a/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh index 9786ba5..653077b 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-1D-global.sh @@ -67,7 +67,8 @@ python -m protossl.protoecgnet_trainer \ 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-classifier/latest/probs.npy \ +--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/probs.npy $RUN_DIR/$EXP_NAME/probs.npy +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-project.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-global-project.sh deleted file mode 100755 index d8bb3d3..0000000 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-global-project.sh +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/bash - -#SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=300gb -#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 -BRANCH="2D-global" -EXP_NAME="$BRANCH-tuned" -# best lambda-sweep trial (wandb sweep f5mqrrme, run vuoyocwi, val_ap=0.8759) -TUNED_TRIAL="v337-27ubgqcf" - -# project in the training dataset -python -m protossl.protoecgnet_trainer \ - --pipeline-stage project-prototypes-supervised \ - --config $REPO_ROOT/configs/$BRANCH.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.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/$TUNED_TRIAL/best.ckpt \ - --data.num_workers 8 \ - --data.prefetch_factor 4 diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh index 6286a53..58211e2 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-global.sh @@ -67,7 +67,8 @@ python -m protossl.protoecgnet_trainer \ 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-classifier/latest/probs.npy \ +--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/probs.npy $RUN_DIR/$EXP_NAME/probs.npy +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-project.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial-project.sh deleted file mode 100755 index ac8019f..0000000 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial-project.sh +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/bash - -#SBATCH --cpus-per-task=24 -#SBATCH --mem-per-gpu=300gb -#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 -BRANCH="2D-partial" -EXP_NAME="$BRANCH-tuned" -# best lambda-sweep trial (wandb sweep blqwqt64, run btq9wfqt, val_ap=0.5352) -TUNED_TRIAL="v294-895m7c3p" - -# project in the training dataset -python -m protossl.protoecgnet_trainer \ - --pipeline-stage project-prototypes-supervised \ - --config $REPO_ROOT/configs/$BRANCH.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.pretrained_weights $RUN_DIR/$EXP_NAME/learn-prototypes-supervised/$TUNED_TRIAL/best.ckpt \ - --data.num_workers 8 \ - --data.prefetch_factor 4 diff --git a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh index bac4280..9f1fba5 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-2D-partial.sh @@ -67,7 +67,8 @@ python -m protossl.protoecgnet_trainer \ 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-classifier/latest/probs.npy \ +--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/probs.npy $RUN_DIR/$EXP_NAME/probs.npy +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.sh b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh index 1812fcb..8a190d3 100755 --- a/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh +++ b/large-user-study/scripts/run-protoecgnet-heedb-fusion.sh @@ -55,7 +55,7 @@ python -m protossl.protoecgnet_trainer \ 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/probs.npy \ +--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/probs.npy $RUN_DIR/$EXP_NAME/probs.npy +cp $RUN_DIR/$EXP_NAME/train-fusion-classifier/latest/test_probs.npy $RUN_DIR/$EXP_NAME/test_probs.npy From e46d7ddb622f456a5c3f9b93b4aa8a40d6cf79a9 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 10 Oct 2026 15:34:25 -0500 Subject: [PATCH 31/33] paper resubmission --- .github/CODEOWNERS | 1 + .gitignore | 1 + README.md | 7 +- protossl/plotting.py | 28 +-- results/ecg-results-rebuttal.ipynb | 27 ++- results/proto-recur-multilabel.ipynb | 256 +++++++++++++++++++++++---- scripts/_submit_job.sh | 2 +- scripts/queue-pclr-logreg.sh | 52 ++++++ user-study/prepare_samples.ipynb | 56 +++--- 9 files changed, 355 insertions(+), 75 deletions(-) create mode 100644 scripts/queue-pclr-logreg.sh diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index d0683b4..21e900c 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -91,6 +91,7 @@ scripts/echonext/run-tabular-logreg.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 diff --git a/.gitignore b/.gitignore index 002c6b0..c21884a 100644 --- a/.gitignore +++ b/.gitignore @@ -221,3 +221,4 @@ results/figs/ ecgfounder-checkpoint/ wandb/ .devcontainer/ +user-study/images-pdf diff --git a/README.md b/README.md index 8b795be..882c42d 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,12 @@ This repo contains the source code for **ProtoSSL**, a novel framework for label ### 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 diff --git a/protossl/plotting.py b/protossl/plotting.py index 59ebaaa..3e20645 100644 --- a/protossl/plotting.py +++ b/protossl/plotting.py @@ -458,8 +458,13 @@ def plot_segment( 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.""" + """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 @@ -478,16 +483,17 @@ def plot_segment( ) # draw lead label to the left of the plot - ax.text( - x=-0.1, - y=0.8, - s=lead, - fontsize=10, - fontweight="bold", - transform=ax.transAxes, - ha="right", - va="bottom", - ) + 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): diff --git a/results/ecg-results-rebuttal.ipynb b/results/ecg-results-rebuttal.ipynb index 5c733dd..f44603d 100644 --- a/results/ecg-results-rebuttal.ipynb +++ b/results/ecg-results-rebuttal.ipynb @@ -582,6 +582,19 @@ "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", @@ -593,7 +606,7 @@ }, { "cell_type": "markdown", - "id": "24", + "id": "25", "metadata": {}, "source": [ "### No Assignment (PIT & PIP) - and No Projection (Latent) (new)" @@ -602,7 +615,7 @@ { "cell_type": "code", "execution_count": null, - "id": "25", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -618,7 +631,7 @@ }, { "cell_type": "markdown", - "id": "26", + "id": "27", "metadata": {}, "source": [ "### ZZU full scale more prototypes in ProtoSSL pretraining" @@ -627,7 +640,7 @@ { "cell_type": "code", "execution_count": null, - "id": "27", + "id": "28", "metadata": {}, "outputs": [], "source": [ @@ -642,7 +655,7 @@ { "cell_type": "code", "execution_count": null, - "id": "28", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -660,7 +673,7 @@ { "cell_type": "code", "execution_count": null, - "id": "29", + "id": "30", "metadata": {}, "outputs": [], "source": [ @@ -681,7 +694,7 @@ { "cell_type": "code", "execution_count": null, - "id": "30", + "id": "31", "metadata": {}, "outputs": [], "source": [] diff --git a/results/proto-recur-multilabel.ipynb b/results/proto-recur-multilabel.ipynb index bcc4bc5..9ea4c68 100644 --- a/results/proto-recur-multilabel.ipynb +++ b/results/proto-recur-multilabel.ipynb @@ -32,6 +32,17 @@ "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", @@ -71,7 +82,7 @@ }, { "cell_type": "markdown", - "id": "2", + "id": "3", "metadata": {}, "source": [ "## Datasets and labels\n", @@ -95,7 +106,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -184,7 +195,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -201,7 +212,7 @@ }, { "cell_type": "markdown", - "id": "5", + "id": "6", "metadata": {}, "source": [ "## Load assignment and projection metadata\n", @@ -215,7 +226,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -249,7 +260,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [], "source": [ @@ -274,7 +285,7 @@ }, { "cell_type": "markdown", - "id": "8", + "id": "9", "metadata": {}, "source": [ "## Prototypes assigned to each label\n", @@ -287,7 +298,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -307,7 +318,7 @@ }, { "cell_type": "markdown", - "id": "10", + "id": "11", "metadata": {}, "source": [ "## Overlap between datasets\n", @@ -321,7 +332,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -360,7 +371,7 @@ { "cell_type": "code", "execution_count": null, - "id": "12", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -370,7 +381,7 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "14", "metadata": {}, "outputs": [], "source": [ @@ -391,7 +402,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14", + "id": "15", "metadata": {}, "outputs": [], "source": [ @@ -412,7 +423,7 @@ { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -428,7 +439,7 @@ { "cell_type": "code", "execution_count": null, - "id": "16", + "id": "17", "metadata": {}, "outputs": [], "source": [ @@ -450,7 +461,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "18", "metadata": {}, "outputs": [], "source": [ @@ -469,30 +480,58 @@ { "cell_type": "code", "execution_count": null, - "id": "18", + "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", - "fig, axs = plt.subplots(n_rows, n_cols, figsize=(5 * n_cols, 4 * n_rows))\n", - "for ax, (label, sets) in zip(axs.flat, PROTOTYPES.items()):\n", - " venn3([sets[name] for name in DATASETS], set_labels=tuple(DATASETS), ax=ax)\n", - " ax.set_title(label)\n", - "for ax in axs.flat[len(PROTOTYPES) :]: # blank out any unused panels\n", - " ax.axis(\"off\")\n", - "fig.tight_layout()\n", - "fig.savefig(FIG_DIR / f\"proto-recur-venn-seed{SEED}.png\", dpi=300, bbox_inches=\"tight\")" + "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": "19", + "id": "20", "metadata": {}, "source": [ "## Recurrent prototypes\n", @@ -504,7 +543,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -523,7 +562,7 @@ }, { "cell_type": "markdown", - "id": "21", + "id": "22", "metadata": {}, "source": [ "## Plot a prototype's projections\n", @@ -536,7 +575,7 @@ { "cell_type": "code", "execution_count": null, - "id": "22", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -565,7 +604,7 @@ { "cell_type": "code", "execution_count": null, - "id": "23", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -604,6 +643,7 @@ " 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", @@ -613,7 +653,7 @@ { "cell_type": "code", "execution_count": null, - "id": "24", + "id": "25", "metadata": {}, "outputs": [], "source": [ @@ -625,7 +665,7 @@ { "cell_type": "code", "execution_count": null, - "id": "25", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -637,7 +677,7 @@ { "cell_type": "code", "execution_count": null, - "id": "26", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -647,7 +687,155 @@ { "cell_type": "code", "execution_count": null, - "id": "27", + "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": [] 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/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/prepare_samples.ipynb b/user-study/prepare_samples.ipynb index 74db6a8..89c12e8 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/labsup-proto-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": [], From 268425e8adb0ebc1ac87391f090ce96cee70cc8a Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 10 Oct 2026 16:22:59 -0500 Subject: [PATCH 32/33] massive cleanup with improved README and extensive rename of experiment paths to align with paper --- .github/CODEOWNERS | 2 +- README.md | 82 +++++++++++--- protossl/protoecgnet_trainer.py | 2 +- protossl/trainer.py | 2 +- pyproject.toml | 2 +- results/audio-results.ipynb | 6 +- results/ecg-results.ipynb | 22 ++-- ...oto-direct.sh => 2-run-supproto-direct.sh} | 2 +- ...pl.sh => 2-z-run-supproto-direct-28ppl.sh} | 2 +- ...ppl.sh => 2-z-run-supproto-direct-7ppl.sh} | 2 +- scripts/3-run-protossl-heedb-pila.sh | 2 +- scripts/3-z-run-protossl-heedb-pila-28ppl.sh | 2 +- scripts/3-z-run-protossl-heedb-pila-7ppl.sh | 2 +- scripts/3-z-run-protossl-heedb-pila-83ppl.sh | 2 +- scripts/4-run-supproto-heedb-rila.sh | 4 +- ...pila.sh => 4-y-run-supproto-heedb-pila.sh} | 4 +- ...h => 4-z-run-supproto-heedb-rila-28ppl.sh} | 4 +- ...sh => 4-z-run-supproto-heedb-rila-7ppl.sh} | 4 +- scripts/6-run-protossl-heedb-pia.sh | 2 +- scripts/6-z-run-protossl-heedb-pia-83ppl.sh | 2 +- scripts/7-1-run-protossl-heedb-pit.sh | 2 +- scripts/7-2-run-protossl-heedb-pip.sh | 2 +- scripts/README.md | 4 +- ...oto-direct.sh => 2-run-supproto-direct.sh} | 2 +- ...ila.sh => 4-run-supproto-audioset-rila.sh} | 2 +- scripts/audio/queue-experiments.sh | 12 +- ...t.sh => voxceleb-2-run-supproto-direct.sh} | 2 +- ... voxceleb-4-run-supproto-audioset-rila.sh} | 2 +- .../ecg-fms/_compute_ecgfounder_embeddings.py | 2 +- scripts/pretrain/run-heedb-normalizations.sh | 2 +- .../run-protossl-heedb-150-pretrain.sh | 4 +- .../run-protossl-heedb-pretrain-w-attn.sh | 6 +- .../pretrain/run-protossl-heedb-pretrain.sh | 6 +- .../run-supproto-heedb-150-pretrain.sh | 4 +- .../pretrain/run-supproto-heedb-pretrain.sh | 6 +- scripts/queue-experiments.sh | 106 +++++++++--------- user-study/prepare_samples.ipynb | 2 +- 37 files changed, 189 insertions(+), 129 deletions(-) rename scripts/{2-run-labsup-proto-direct.sh => 2-run-supproto-direct.sh} (98%) rename scripts/{2-z-run-labsup-proto-direct-28ppl.sh => 2-z-run-supproto-direct-28ppl.sh} (98%) rename scripts/{2-z-run-labsup-proto-direct-7ppl.sh => 2-z-run-supproto-direct-7ppl.sh} (98%) rename scripts/{4-y-run-labsup-proto-heedb-pila.sh => 4-y-run-supproto-heedb-pila.sh} (97%) rename scripts/{4-z-run-labsup-proto-heedb-rila-28ppl.sh => 4-z-run-supproto-heedb-rila-28ppl.sh} (97%) rename scripts/{4-z-run-labsup-proto-heedb-rila-7ppl.sh => 4-z-run-supproto-heedb-rila-7ppl.sh} (97%) rename scripts/audio/{2-run-labsup-proto-direct.sh => 2-run-supproto-direct.sh} (98%) rename scripts/audio/{4-run-labsup-proto-audioset-rila.sh => 4-run-supproto-audioset-rila.sh} (99%) rename scripts/audio/{voxceleb-2-run-labsup-proto-direct.sh => voxceleb-2-run-supproto-direct.sh} (98%) rename scripts/audio/{voxceleb-4-run-labsup-proto-audioset-rila.sh => voxceleb-4-run-supproto-audioset-rila.sh} (99%) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 21e900c..2df946a 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -69,7 +69,7 @@ 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-* @StevenSong scripts/4-* @StevenSong diff --git a/README.md b/README.md index 882c42d..8b9f929 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# 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 ``` @@ -12,14 +12,23 @@ This repo contains the source code for **ProtoSSL**, a novel framework for label } ``` -## 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 @@ -27,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 @@ -52,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/protossl/protoecgnet_trainer.py b/protossl/protoecgnet_trainer.py index c2db0d4..675989c 100644 --- a/protossl/protoecgnet_trainer.py +++ b/protossl/protoecgnet_trainer.py @@ -702,7 +702,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}" ) diff --git a/protossl/trainer.py b/protossl/trainer.py index e2016c6..db047ab 100644 --- a/protossl/trainer.py +++ b/protossl/trainer.py @@ -1035,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}" ) 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/results/audio-results.ipynb b/results/audio-results.ipynb index 01b6815..bc0e5b0 100644 --- a/results/audio-results.ipynb +++ b/results/audio-results.ipynb @@ -21,9 +21,9 @@ "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", diff --git a/results/ecg-results.ipynb b/results/ecg-results.ipynb index 2ec018b..6399122 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", 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-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-rila.sh b/scripts/4-run-supproto-heedb-rila.sh index 3c49364..1b30aac 100755 --- a/scripts/4-run-supproto-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/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-1-run-protossl-heedb-pit.sh b/scripts/7-1-run-protossl-heedb-pit.sh index 0035093..acefad5 100755 --- a/scripts/7-1-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/7-2-run-protossl-heedb-pip.sh b/scripts/7-2-run-protossl-heedb-pip.sh index ea25fee..a9e84c8 100755 --- a/scripts/7-2-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/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/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/4-run-labsup-proto-audioset-rila.sh b/scripts/audio/4-run-supproto-audioset-rila.sh similarity index 99% rename from scripts/audio/4-run-labsup-proto-audioset-rila.sh rename to scripts/audio/4-run-supproto-audioset-rila.sh index a2e578a..3b01623 100755 --- a/scripts/audio/4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/4-run-supproto-audioset-rila.sh @@ -27,7 +27,7 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts/audio # experiment parameters -EXP_NAME="labsup-proto-audioset-rila" +EXP_NAME="supproto-audioset-rila" PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2") python -m protossl.trainer \ diff --git a/scripts/audio/queue-experiments.sh b/scripts/audio/queue-experiments.sh index 6b746c1..8f65d2a 100644 --- a/scripts/audio/queue-experiments.sh +++ b/scripts/audio/queue-experiments.sh @@ -1,17 +1,17 @@ #!/bin/bash scripts=( - # "1-run-blackbox-direct.sh" - # "2-run-labsup-proto-direct.sh" + "1-run-blackbox-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" + "1-run-blackbox-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 diff --git a/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh b/scripts/audio/voxceleb-2-run-supproto-direct.sh similarity index 98% rename from scripts/audio/voxceleb-2-run-labsup-proto-direct.sh rename to scripts/audio/voxceleb-2-run-supproto-direct.sh index 4f70c6d..272973a 100755 --- a/scripts/audio/voxceleb-2-run-labsup-proto-direct.sh +++ b/scripts/audio/voxceleb-2-run-supproto-direct.sh @@ -31,7 +31,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/voxceleb-4-run-labsup-proto-audioset-rila.sh b/scripts/audio/voxceleb-4-run-supproto-audioset-rila.sh similarity index 99% rename from scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh rename to scripts/audio/voxceleb-4-run-supproto-audioset-rila.sh index d3b0d92..77e7e88 100755 --- a/scripts/audio/voxceleb-4-run-labsup-proto-audioset-rila.sh +++ b/scripts/audio/voxceleb-4-run-supproto-audioset-rila.sh @@ -31,7 +31,7 @@ echo "Using REPO_ROOT=$REPO_ROOT" cd $REPO_ROOT/scripts/audio # experiment parameters -EXP_NAME="labsup-proto-audioset-rila" +EXP_NAME="supproto-audioset-rila" PRETRAIN_RUN=$(realpath -m "$RUN_DIR/../../supproto-audioset-v2") python -m protossl.trainer \ 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/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-protossl-heedb-150-pretrain.sh b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh index 09e8db3..0f99d2a 100755 --- a/scripts/pretrain/run-protossl-heedb-150-pretrain.sh +++ b/scripts/pretrain/run-protossl-heedb-150-pretrain.sh @@ -12,8 +12,8 @@ set -e PRETRAIN_DATASET=/opt/gpudata/ecg/heedb -RUN_DIR=/home/songs1/protossl-ecg-outputs-rebuttal -REPO_ROOT=/home/songs1/ProtoSSL +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}" diff --git a/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh b/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.sh index 4841c16..ae09c73 100755 --- a/scripts/pretrain/run-protossl-heedb-pretrain-w-attn.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-w-attn" +EXP_NAME="protossl-heedb-w-attn" BACKBONE=resnet18 CONV=2D diff --git a/scripts/pretrain/run-protossl-heedb-pretrain.sh b/scripts/pretrain/run-protossl-heedb-pretrain.sh index 2255794..a05532f 100755 --- a/scripts/pretrain/run-protossl-heedb-pretrain.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-pretrain.sh b/scripts/pretrain/run-supproto-heedb-150-pretrain.sh index 762717c..f71c313 100755 --- a/scripts/pretrain/run-supproto-heedb-150-pretrain.sh +++ b/scripts/pretrain/run-supproto-heedb-150-pretrain.sh @@ -12,8 +12,8 @@ set -e PRETRAIN_DATASET=/opt/gpudata/ecg/heedb -RUN_DIR=/home/songs1/protossl-ecg-outputs-rebuttal -REPO_ROOT=/home/songs1/ProtoSSL +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}" diff --git a/scripts/pretrain/run-supproto-heedb-pretrain.sh b/scripts/pretrain/run-supproto-heedb-pretrain.sh index 44833e9..931a24e 100755 --- a/scripts/pretrain/run-supproto-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}" @@ -19,7 +19,7 @@ cd $REPO_ROOT/scripts 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 e947cd2..5d60e5c 100755 --- a/scripts/queue-experiments.sh +++ b/scripts/queue-experiments.sh @@ -2,8 +2,7 @@ set -e -declare -a DATASETS=("zzu") -# declare -a DATASETS=("echonext" "ptbxl" "cinc" "mimic" "zzu" "code15") +declare -a DATASETS=("echonext" "ptbxl" "cinc" "mimic" "zzu" "code15") declare -a SEEDS=(42 67 70 73 99) declare -A DATASET_DIRS=( ["echonext"]="echonext" @@ -18,8 +17,7 @@ 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=("") -# declare -a SUFFIXES_zzu=("" "-4k" "-2k" "-1k" "-512" "-256") +declare -a SUFFIXES_zzu=("" "-4k" "-2k" "-1k" "-512" "-256") source _submit_job.sh @@ -44,53 +42,59 @@ for seed in "${SEEDS[@]}"; do for suffix in "${suffixes[@]}"; do echo "SUFFIX: $suffix" - # cache_id=$(submit_job "$suffix" 0-run-cache-data.sh) - # 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" 3-run-protossl-heedb-pila.sh "--dependency=afterok:$cache_id" - # submit_job "$suffix" 4-run-supproto-heedb-rila.sh "--dependency=afterok:$cache_id" - submit_job "$suffix" 3-run-protossl-heedb-150-pila.sh - # 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" - - # 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" == "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" 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" - - # # supproto no-proj - # submit_job "$suffix" 4-y-run-labsup-proto-heedb-pila.sh "--dependency=afterok:$cache_id" - - # # assignment method - # submit_job "$suffix" 6-run-protossl-heedb-pia.sh "--dependency=afterok:$cache_id" - - # # 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-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" - # submit_job "$suffix" 7-3-run-protossl-heedb-latent.sh - - # # start with pretrained encoder - # patches_id=$(submit_job "$suffix" 9-0-run-ecgfounder-patches.sh) - # echo $patches_id - # submit_job "$suffix" 9-1-run-ecgfounder-lap.sh "--dependency=afterok:$patches_id" - # submit_job "$suffix" 9-2-run-ecgfounder-clustering.sh "--dependency=afterok:$patches_id" - # submit_job "$suffix" 9-3-run-ecgfounder-random.sh "--dependency=afterok:$patches_id" - # fi + cache_id=$(submit_job "$suffix" 0-run-cache-data.sh) + echo $cache_id + + submit_job "$suffix" 1-run-blackbox-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-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-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-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-supproto-heedb-pila.sh "--dependency=afterok:$cache_id" + + # assignment method + submit_job "$suffix" 6-run-protossl-heedb-pia.sh "--dependency=afterok:$cache_id" + + # 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-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) + echo $patches_id + submit_job "$suffix" 9-1-run-ecgfounder-lap.sh "--dependency=afterok:$patches_id" + submit_job "$suffix" 9-2-run-ecgfounder-clustering.sh "--dependency=afterok:$patches_id" + submit_job "$suffix" 9-3-run-ecgfounder-random.sh "--dependency=afterok:$patches_id" + fi done done done diff --git a/user-study/prepare_samples.ipynb b/user-study/prepare_samples.ipynb index 89c12e8..4740f9b 100644 --- a/user-study/prepare_samples.ipynb +++ b/user-study/prepare_samples.ipynb @@ -55,7 +55,7 @@ "labels = [\"AMI\", \"CRBBB\", \"CLBBB\", \"PVC\"]\n", "cases_per_label = 5 # CPL\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/labsup-proto-direct\")\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", From e391199752023a416978721901fcce857dfeab05 Mon Sep 17 00:00:00 2001 From: StevenSong Date: Sat, 10 Oct 2026 16:46:10 -0500 Subject: [PATCH 33/33] cleanup result notebooks --- results/audio-results.ipynb | 6 +++--- results/ecg-results.ipynb | 35 ++--------------------------------- results/meta-results.ipynb | 12 ++++++------ 3 files changed, 11 insertions(+), 42 deletions(-) diff --git a/results/audio-results.ipynb b/results/audio-results.ipynb index bc0e5b0..10e25bc 100644 --- a/results/audio-results.ipynb +++ b/results/audio-results.ipynb @@ -30,9 +30,9 @@ "\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.ipynb b/results/ecg-results.ipynb index 6399122..1307a48 100644 --- a/results/ecg-results.ipynb +++ b/results/ecg-results.ipynb @@ -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",