copy data from repo

README.md

@@ -1,3 +1,85 @@
----
-license: afl-3.0
----
+# Kardio-Net: A Deep Learning Model for Predicting Serum Potassium Levels from Apple Watch ECG in ESRD Patients
+
+
+### I. Input preparation
+
+
+1. Prerequisites
+
+    Ensure that all ECG files are saved in *.npy format with the following shapes:
+    - (5000, 12) for 12-lead ECGs
+    - (500 * 30 seconds, 1) for Apple Watch ECGs
+    
+    Store these files in a single flat directory, and include a manifest CSV file in the same location to accompany 
+    them.
+    
+
+2. Manifest File Format
+
+    The manifest CSV file should include a header.
+    Each row corresponds to one ECG file with the following columns:
+
+    - filename: Name of the .npy file (without the extension).
+    - label: serum potassium label.
+
+
+
+### II. Inference
+
+
+### For 12-Lead ECG
+
+<!-- #region -->
+1. use predict_potassium_12lead.py to get the potassium level prediction.
+
+
+2. Edit data_path and manifest_path and run the predict_potassium.py script.
+
+
+3. Upon completion, a file named "dataloader_0_predictions.csv" will be saved in the same directory. This file contains the inference results "preds" from the model.
+
+
+4. Use generate_result.py to get the performance metric and figure.
+<!-- #endregion -->
+
+### For Single Lead ECG
+
+
+### A. ECG preprocessing and segmentation
+
+<!-- #region -->
+1. Use preprocessing.py for denoise, normalize, and segment ECG into 5-second for input
+
+
+2. Set the following paths in preprocessing.py:
+
+    - raw_directory = "path/to/raw_data_directory"  #raw ecg folder for target task
+    - output_directory = "path/to/output_directory" #output folder for normalize ECG
+    - manifest_path = "/path/to/manifest.csv"  # Manifest file path
+    - output_path = "path/to/output_path"  # Output path for segmented ECGs
+
+
+3. Execute predict.py.
+
+
+4. If the ECG files were already normalize, can execute the segmentation function only.
+<!-- #endregion -->
+
+### B. potassium regression model
+
+<!-- #region -->
+1. use predict_potassium_1lead.py to get the potassium level prediction.
+
+
+2. Edit data_path and manifest_path and run the predict_potassium.py script.
+
+
+3. Upon completion, a file named "dataloader_0_predictions.csv" will be saved in the same directory. This file contains the inference results "preds" from the model.
+
+
+4. Use generate_result.py to get the performance metric and figure.
+<!-- #endregion -->
+
+```python
+
+```

noise_classifier.py

@@ -0,0 +1,72 @@
+# An example of how to run inference using a model trained with cvair.
+# We want to use a pretrained model to make predictions on a dataset of new examples.
+
+import numpy as np
+import pandas as pd
+import torch
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+from utils.datasets import ECGSingleLeadDataset
+from utils.models import EffNet
+from utils.training_models import BinaryClassificationModel
+
+
+def append_noise_predictions_to_manifest(data_path, manifest_path, weights_path):
+    # Initialize a dataset
+    test_ds = ECGSingleLeadDataset(
+        data_path=data_path,
+        manifest_path=manifest_path,
+        update_manifest_func=None,
+    )
+
+    # Wrap the dataset in a dataloader
+    test_dl = DataLoader(
+        test_ds, 
+        num_workers=16, 
+        batch_size=512, 
+        drop_last=False, 
+        shuffle=False
+    )
+
+    # Initialize the backbone model
+    backbone = EffNet(input_channels=1, output_neurons=1)
+
+    # Pass the backbone to a wrapper
+    model = BinaryClassificationModel(backbone)
+
+    # Load the pretrained weights
+    weights = torch.load(weights_path)
+    model.load_state_dict(weights)
+
+    # Initialize a Trainer object
+    trainer = Trainer(accelerator="gpu", devices=1)
+
+    # Run inference
+    trainer.predict(model, dataloaders=test_dl)
+
+    # Read the predictions CSV file
+    df = pd.read_csv('dataloader_0_predictions.csv')
+
+    # Normalize predictions
+    max_preds = df['preds'].max()
+    min_preds = df['preds'].min()
+    df['noise_preds_normal'] = (df['preds'] - min_preds) / (max_preds - min_preds)
+
+    # Drop the original predictions column
+    df.drop(columns='preds', inplace=True)
+
+    # Save the modified dataframe to the original manifest path
+    df.to_csv(manifest_path)
+
+
+
+if __name__ == "__main__":
+
+    data_path="/your/wearable/ecg/data path/"
+
+    manifest_path="manifest"
+    
+    weights_path = "model_noise_classifier.pt"
+    
+    append_noise_predictions_to_manifest(data_path, manifest_path, weights_path)
+

predict_potassium_12lead.py

@@ -0,0 +1,46 @@
+import torch
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+from utils.datasets import ECGDataset
+from utils.models import EffNet
+from utils.training_models import RegressionModel
+
+
+# +
+# This is the path where your data samples are stored.
+data_path = "your/ecg/data/folder"
+
+# This is the path where your manifest, containing filenames for inference to be run on, is stored.
+manifest_path = 'your/manifest/path'
+# -
+
+
+# Initialize a dataset that contains the examples you want to run prediction on. 
+test_ds = ECGDataset(
+    split="test",
+    data_path=data_path,
+    manifest_path=manifest_path,
+    update_manifest_func=None,
+)
+
+# Wrap the dataset in a dataloader to handle batching and multithreading.
+test_dl = DataLoader(
+    test_ds, 
+    num_workers=16, 
+    batch_size=256, 
+    drop_last=False, 
+    shuffle=False
+)
+
+# Initialize the "backbone", the core model weights that will act on the data.
+backbone = EffNet(input_channels=12, output_neurons=1)
+
+model = RegressionModel(backbone)
+
+weights = torch.load("model_12_lead.pt")
+print(model.load_state_dict(weights))
+
+# +
+trainer = Trainer(accelerator="gpu", devices=1)
+
+trainer.predict(model, dataloaders=test_dl)

predict_potassium_1lead.py

@@ -0,0 +1,44 @@
+import torch
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+from utils.datasets import ECGSingleLeadDataset
+from utils.models import EffNet
+from utils.training_models import RegressionModel
+
+# +
+# This is the path where your data samples are stored.
+data_path = "your/ecg/data/folder"
+
+# This is the path where your manifest, containing filenames for inference to be run on, is stored.
+manifest_path = 'your/manifest/path'
+# -
+
+# Initialize a dataset that contains the examples you want to run prediction on.
+test_ds = ECGSingleLeadDataset(
+    data_path=data_path,
+    manifest_path=manifest_path,
+    update_manifest_func=None,
+)
+
+# Wrap the dataset in a dataloader to handle batching and multithreading.
+test_dl = DataLoader(
+    test_ds, 
+    num_workers=16, 
+    batch_size=512, 
+    drop_last=False, 
+    shuffle=False
+)
+
+# +
+backbone = EffNet()
+
+model = RegressionModel(backbone)
+# -
+
+weights = torch.load("model_single_lead_5seconds_length.pt")
+print(model.load_state_dict(weights))
+
+# +
+trainer = Trainer(accelerator="gpu", devices=1)
+
+trainer.predict(model, dataloaders=test_dl)

preprocessing.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python
+# coding: utf-8
+# %%
+import numpy as np
+import os
+from pathlib import Path
+from tqdm import tqdm
+from utils.ecg_utils import (
+    remove_baseline_wander,
+    wavelet_denoise_signal,
+    plot_12_lead_ecg,
+)
+
+
+# %%
+def calculate_means_stds(npy_directory, n):
+    
+    npy_directory = Path(npy_directory)
+    filelist = os.listdir(npy_directory)
+    np.random.shuffle(filelist)
+
+    full_batch = np.zeros((n, 5000, 12))
+    count = 0
+
+    for i, npy_filename in enumerate(tqdm(filelist[:n])):
+        npy_filepath = npy_directory / npy_filename
+        ekg_numpy_array = np.load(npy_filepath)
+
+        if ekg_numpy_array.shape[0] != 5000:
+            continue
+
+        full_batch[count] = ekg_numpy_array
+        count += 1
+
+    full_batch = full_batch[:count]  # Trim the array to remove unused entries
+    ecg_means = np.mean(full_batch, axis=(0, 1))
+    ecg_stds = np.std(full_batch, axis=(0, 1))
+
+    if ecg_means.shape[0] == ecg_stds.shape[0] == 12:
+        print('Shape of mean and std for ECG normalization are correct!')
+
+    return ecg_means, ecg_stds
+
+
+# %%
+# run the function on a list of filenames.
+def ecg_denoising(
+        raw_directory=raw_directory,
+        output_directory=output_directory,
+        ecg_means = ecg_means,
+        ecg_stds = ecg_stds
+    ):
+
+    filelist = os.listdir(raw_directory)
+    
+    for i, filename in enumerate(tqdm(filelist[:n])):
+        
+        # Signal processing
+        raw_directory = Path(raw_directory)
+        ecg_filepath = raw_directory / filename
+        ecg_numpy_array = np.load(ecg_filepath)
+        # 1. Wandering baseline removal
+        ecg_numpy_array = remove_baseline_wander(
+            ecg_numpy_array, sampling_frequency=sampling_frequency
+        )
+
+        # Discrete wavelet transform denoising
+        for lead in range(12):
+            ecg_numpy_array[:, lead] = wavelet_denoise_signal(ecg_numpy_array[:, lead])
+
+        # Lead-wise normalization with precomputed means and standard deviations
+        ecg_numpy_array = (ecg_numpy_array - ecg_means) / ecg_stds
+
+        np.save(output_directory / filename, ecg_numpy_array)
+        
+        return True
+
+
+# %%
+def segmentation(data_path, manifest_path, output_path, length, steps):
+    manifest = pd.read_csv(manifest_path)
+
+    data = []
+    print('Staring segmenting ECG......')
+    for index in tqdm(range(manifest.shape[0])):
+        mrn = manifest['MRN'].iloc[index]            #MRN as column name for medical record number
+        filename = manifest['filename'].iloc[index]  #filename as column name for ecg filename
+        k = manifest['TEST_RSLT'].iloc[index]        #TEST_RSLT as column name for potassium level
+
+        ecg_array = np.load(os.path.join(data_path, filename))
+        ecg_array = ecg_array[:, 0]  # assume lead I is the first lead in npy file
+
+        # Loop through every second as start point:
+        for start in range(0, len(ecg_array), 500*steps):
+            end = start + 500 * length    # 500 points for each seconds
+
+            if start >= 0 and end <= len(ecg_array):
+                sample = ecg_array[start:end]
+
+                if len(sample) == 500 * length:
+                    data.append({'mrn': mrn, 'original_filename': filename, 'ecg': sample, 'label': k})
+                else:
+                    print(f'Different sample size for {filename}: {len(sample)}')
+
+    df = pd.DataFrame(data)
+    df['filename'] = None
+
+    if not os.path.exists(output_path):
+        os.makedirs(output_path)
+    
+    print('Saving segmented ECG......')
+    for index, row in df.iterrows():
+        original_filename = row['original_filename']
+        ecg_array = row['ecg']
+        new_file_name = f"{original_filename.rsplit('.', 1)[0]}_{index+1}.npy"
+        new_file_path = os.path.join(output_path, new_file_name)
+        
+        df.at[index, 'filename'] = new_file_name
+        np.save(new_file_path, ecg_array)
+
+    df.drop(columns=['ecg'], inplace=True)
+    df.to_csv(f'{output_path}/{length}seconds_length_{steps}seconds_step_ecg_manifest.csv')
+
+
+# %%
+if __name__ == "__main__":
+    npy_directory = "ecg folder for entire database" 
+    n = 100000  # the number of ecg for calculating mean and std
+    
+    print('Calculating ECG means and stds........')
+    ecg_means, ecg_stds = calculate_means_stds(npy_directory, n)
+
+    raw_directory = "path/to/raw_data_directory"  #raw ecg folder for target task
+    output_directory = "path/to/output_directory" #output ecg folder for target task
+
+    print('Denoising and Normalizing ECGs........')
+    ecg_denoising(raw_directory, output_directory, ecg_means, ecg_stds)
+    
+    
+    data_path = output_directory  # Output directory from the above step
+    manifest_path = "/path/to/manifest.csv"  # Manifest file path
+    output_path = "path/to/output_path"  # Output path for segmented ECGs
+    length = 5  # Length of each segment in seconds
+    steps = 1  # Number of seconds step
+
+    process_ecg(data_path, manifest_path, output_path, length, steps)
+

utils/datasets.py

@@ -0,0 +1,392 @@
+import os
+import random
+from pathlib import Path
+from typing import Callable, List, Tuple, Optional, Iterable, Dict, Union
+from typing_extensions import TypedDict, Unpack, Required, NotRequired
+
+import numpy as np
+import pandas as pd
+import torch
+import torchvision.transforms.functional as TF
+import wandb
+from torch.utils.data import Dataset
+import torch.nn as nn
+
+
+def format_mrn(mrn):
+    return str(mrn).strip().zfill(20)
+
+
+class CedarsDatasetTypeAnnotations(TypedDict, total=False):
+    """A dummy class used to make IDE autocomplete and tooltips work properly with how we pass **kwargs through in subclasses of CedarsDataset."""
+    data_path: Required[Union[Path, str]]
+    manifest_path: Required[Union[Path, str]]
+    split: NotRequired[str]
+    labels: NotRequired[Iterable[str]]
+    extra_inputs: NotRequired[Iterable[str]]
+    update_manifest_func: NotRequired[Callable[[pd.DataFrame], pd.DataFrame]]
+    subsample: NotRequired[Union[Path, str]]
+    augmentations: NotRequired[Union[Iterable[Callable[[torch.Tensor], torch.Tensor]], Callable[[dict], dict], nn.Module]]
+    apply_augmentations_to: NotRequired[Iterable[str]]
+    verify_existing: NotRequired[bool]
+    drop_na_labels: NotRequired[bool]
+    verbose: NotRequired[bool]
+
+
+class CedarsDataset(Dataset):
+    """
+    Generic parent class for several differnet kinds of common datasets we use here at Cedars CVAIR.
+
+    Expects to be used in a scenario where you have a big folder full of input examples (videos, ecgs, 3d arrays, images, etc.) and a big CSV that contains metadata and labels for those examples, called a 'manifest'.
+
+    Args:
+        data_path: Path to a directory full of files you want the dataset to load from.
+        manifest_path: Path to a CSV or Parquet file containing the names, labels, and/or metadata of your files.
+        split: Optional. Allows user to select which split of the manifest to use, assuming the presence of a categorical 'split' column. Defaults to None, meaning that the entire manifest is used by default.
+        extra_inputs: Optional. A list of column names in the manifest that contain additional inputs to the model. Defaults to None.
+        labels: Optional. Name(s) of column(s) in your manifest which contain training labels, in the order you want them returned. If set to None, the dataset will not return any labels, only filenames and inputs. Defaults to None.
+        update_manifest_func: Optional. Allows user to pass in a function to preprocess the manifest after it is loaded, but before the dataset does anything to it.
+        subsample: Optional. A number indicating how many examples to randomly subsample from the manifest. Defaults to None.
+        verbose: Whether to print out progress statements when initializing. Defaults to True.
+        augmentations: Optional. Can be a list of augmentation functions which take in a tensor and return a tensor, a single custom augmentation function which takes in a dict and returns a dict, or a single nn.Module. Defaults to None.
+        apply_augmentations_to: Optional. A list of strings indicating which batch elements to apply augmentations to. Defaults to ("primary_input").
+    """
+
+    def __init__(
+        self,
+        data_path,
+        manifest_path=None,
+        split=None,
+        labels=None,
+        extra_inputs=None,
+        update_manifest_func=None,
+        subsample=None,
+        augmentations=None,
+        apply_augmentations_to=("primary_input",),
+        verify_existing=True,
+        drop_na_labels=True,
+        verbose=True,
+    ):
+
+        self.data_path = Path(data_path)
+        self.augmentations = augmentations
+        self.apply_augmentations_to = apply_augmentations_to
+        self.extra_inputs = extra_inputs
+        self.labels = labels
+
+        if isinstance(self.augmentations, nn.Module):
+            self.augmentations = [self.augmentations]
+
+        if (self.labels is None) and verbose:
+            print(
+                "No label column names were provided, only filenames and inputs will be returned."
+            )
+        if (self.labels is not None) and isinstance(self.labels, str):
+            self.labels = [self.labels]
+        if (self.extra_inputs is not None) and isinstance(self.extra_inputs, str):
+            self.extra_inputs = [self.extra_inputs]
+
+        # Read manifest file
+        if manifest_path is not None:
+            self.manifest_path = Path(manifest_path)
+        else:
+            self.manifest_path = self.data_path / "manifest.csv"
+
+        if self.manifest_path.exists():
+            if self.manifest_path.suffix == ".csv":
+                self.manifest = pd.read_csv(self.manifest_path, low_memory=False)
+            elif self.manifest_path.suffix == ".parquet":
+                self.manifest = pd.read_parquet(self.manifest_path)
+        else:
+            self.manifest = pd.DataFrame(
+                {
+                    "filename": os.listdir(self.data_path),
+                }
+            )
+
+        # do manifest processing that's specific to a given task (different from update_manifest_func,
+        # exists as a method overridden in child classes)
+        self.manifest = self.process_manifest(self.manifest)
+
+        # Apply user-provided update function to manifest
+        if update_manifest_func is not None:
+            self.manifest = update_manifest_func(self, self.manifest)
+
+        # Usually set to "train", "val", or "test". If set to None, the entire manifest is used.
+        if split is not None:
+            self.manifest = self.manifest[self.manifest["split"] == split]
+        if verbose:
+            print(
+                f"Manifest loaded. \nSplit: {split}\nLength: {len(self.manifest):,}"
+            )
+
+        # Make sure all files actually exist. This can be disabled for efficiency if
+        # you have an especially large dataset
+        if verify_existing and "filename" in self.manifest:
+            old_len = len(self.manifest)
+            existing_files = os.listdir(self.data_path)
+            self.manifest = self.manifest[
+                self.manifest["filename"].isin(existing_files)
+            ]
+            new_len = len(self.manifest)
+            if verbose:
+                print(
+                    f"{old_len - new_len} files in the manifest are missing from {self.data_path}."
+                )
+        elif (not verify_existing) and verbose:
+            print(
+                f"self.verify_existing is set to False, so it's possible for the manifest to contain filenames which are not present in {data_path}"
+            )
+
+        # Option to subsample dataset for doing smaller, faster runs
+        if subsample is not None:
+            if isinstance(subsample, int):
+                self.manifest = self.manifest.sample(n=subsample)
+            else:
+                self.manifest = self.manifest.sample(frac=subsample)
+            if verbose:
+                print(f"{subsample} examples subsampled.")
+
+        # Make sure that there are no NAN labels
+        if (self.labels is not None) and drop_na_labels:
+            old_len = len(self.manifest)
+            self.manifest = self.manifest.dropna(subset=self.labels)
+            new_len = len(self.manifest)
+            if verbose:
+                print(
+                    f"{old_len - new_len} examples contained NaN value(s) in their labels and were dropped."
+                )
+        elif (self.labels is not None) and (not drop_na_labels):
+            print(
+                "drop_na_labels is set to False, so it's possible for the manifest to contain NaN-valued labels."
+            )
+
+        # Save manifest to weights and biases run directory
+        if wandb.run is not None:
+            run_data_path = Path(wandb.run.dir).parent / "data"
+            if not run_data_path.is_dir():
+                run_data_path.mkdir()
+
+            save_name = "manifest.csv"
+            if split is not None:
+                save_name = f"{split}_{save_name}"
+
+            self.manifest.to_csv(run_data_path / save_name)
+
+            if verbose:
+                print(f"Copy of manifest saved to {run_data_path}")
+
+    def __len__(self) -> int:
+        return len(self.manifest)
+
+    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
+        output = {}
+        row = self.manifest.iloc[index]
+        if "filename" in row:
+            output["filename"] = row["filename"]
+        if self.labels is not None:
+            output["labels"] = torch.FloatTensor(row[self.labels])
+        file_results = self.read_file(self.data_path / output["filename"], row)
+        if isinstance(file_results, dict):
+            output.update(file_results)
+        else:
+            output["primary_input"] = file_results
+
+        if self.extra_inputs is not None:
+            output["extra_inputs"] = row["extra_inputs"]
+
+        if self.augmentations is not None:
+            output = self.augment(output)
+
+        return output
+
+    def process_manifest(self, manifest: pd.DataFrame) -> pd.DataFrame:
+        if "mrn" in manifest.columns:
+            manifest["mrn"] = manifest["mrn"].apply(format_mrn)
+        if "study_date" in manifest.columns:
+            manifest["study_date"] = pd.to_datetime(manifest["study_date"])
+        if "dob" in manifest.columns:
+            manifest["dob"] = pd.to_datetime(
+                manifest["dob"], infer_datetime_format=True, errors="coerce"
+            )
+        if ("study_date" in manifest.columns) and ("dob" in manifest.columns):
+            manifest["study_age"] = (
+                manifest["study_date"] - manifest["dob"]
+            ) / np.timedelta64(1, "Y")
+        return manifest
+
+    def augment(self, output_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+
+        if isinstance(self.augmentations, Iterable):
+            # would use torch.stack here for cleanliness, but it seems that torchvision
+            # transforms v1's claims about supporting "arbitrary leading dimensions" is
+            # hogwash. they only support up to 4D. so we have to concatenate along the
+            # channel dimension, then apply the augmentations, then split along the channel
+            # dimension.
+            augmentable_inputs = torch.cat(
+                [output_dict[key] for key in self.apply_augmentations_to], dim=0
+            )  # (C*N, T, H, W)
+
+            for aug in self.augmentations:
+                augmentable_inputs = aug(augmentable_inputs)
+
+            place = 0
+            for i, key in enumerate(self.apply_augmentations_to):
+                n_channels = output_dict[key].shape[0]
+                output_dict[key] = augmentable_inputs[place:place+n_channels]
+                place += n_channels
+
+        elif isinstance(self.augmentations, Callable):
+            output_dict = self.augmentations(output_dict)
+
+        else:
+            raise Exception(
+                "self.augmentations must be either an Iterable of augmentations or a single custom augmentation function."
+            )
+
+        return output_dict
+
+    def read_file(self, filepath: Path, row: Optional[pd.Series] = None) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class ECGDataset(CedarsDataset):
+    def __init__(
+        self,
+        # CedarsDataset params
+        data_path: Union[Path, str],
+        manifest_path: Union[Path, str] = None,
+        split: str = None,
+        labels: Union[List[str], str] = None,
+        update_manifest_func: Callable = None,
+        subsample: float = None,
+        verbose: bool = True,
+        verify_existing: bool = True,
+        drop_na_labels: bool = True,
+        # ECGoDataset params
+        leads: List[str] = None,
+        random_lead: bool = False,  # New parameter for random lead selection
+        data_length: int = 5000,
+        **kwargs,
+    ):
+        """
+        Args:
+            leads: List[str] -- which leads you want passed to the model. Defaults to all 12.
+        """
+
+        super().__init__(
+            data_path=data_path,
+            manifest_path=manifest_path,
+            split=split,
+            labels=labels,
+            update_manifest_func=update_manifest_func,
+            subsample=subsample,
+            verbose=verbose,
+            verify_existing=verify_existing,
+            drop_na_labels=drop_na_labels,
+            **kwargs,
+        )
+
+        self.lead_order = [
+            "I",
+            "II",
+            "III",
+            "aVR",
+            "aVL",
+            "aVF",
+            "V1",
+            "V2",
+            "V3",
+            "V4",
+            "V5",
+            "V6",
+        ]
+        self.leads = leads
+        if self.leads is None:
+            self.leads = self.lead_order
+        if isinstance(self.leads, str):
+            self.leads = [self.leads]
+
+        if "first_lead_only" in kwargs:
+            raise (
+                Exception(
+                    '"first_lead_only" has been deprecated. Please pass leads=["I"] \
+                    instead if you would like to train on only the first lead.'
+                )
+            )
+        
+        self.random_lead = random_lead  # Storing the random_lead attribute
+        
+        self.data_length = data_length
+
+
+    def read_file(self, filepath, row=None):
+        # ECGs are usually stored as .npy files.
+        file = np.load(filepath)
+        if file.shape[0] != 12:
+            file = file.T
+        file = torch.tensor(file).float()
+        
+        # Slice the data to the specified length
+        file = file[:, :self.data_length]
+        
+        if self.random_lead:
+            lead_idx = random.choice(range(12))
+            file = file[lead_idx:lead_idx+1]  # Select the random lead
+        else:
+            channels = [self.lead_order.index(lead) for lead in self.leads]
+            file = file[channels]
+
+
+        # Final shape should ideally be NumLeadsxTime(or NumLeadsxTime depending on the resolution of the ECG)
+        return file
+
+
+class ECGSingleLeadDataset(CedarsDataset):
+    def __init__(
+        self,
+        # CedarsDataset params
+        data_path: Union[Path, str],
+        manifest_path: Union[Path, str] = None,
+        labels: Union[List[str], str] = None,
+        update_manifest_func: Callable = None,
+        subsample: float = None,
+        verbose: bool = True,
+        verify_existing: bool = True,
+        drop_na_labels: bool = True,
+        **kwargs,
+    ):
+        """
+        Args:
+            leads: List[str] -- which leads you want passed to the model. Defaults to all 12.
+        """
+
+        super().__init__(
+            data_path=data_path,
+            manifest_path=manifest_path,
+            labels=labels,
+            update_manifest_func=update_manifest_func,
+            subsample=subsample,
+            verbose=verbose,
+            verify_existing=verify_existing,
+            drop_na_labels=drop_na_labels,
+            **kwargs,
+        )
+
+
+    def read_file(self, filepath, row=None):
+        # ECGs are usually stored as .npy files.
+        try:
+            file = np.load(filepath)
+        except Exception as e:
+            print(filepath)
+            print(e)
+
+        file = torch.tensor(file).float().unsqueeze(0)
+
+
+        # Final shape should ideally be NumLeadsxTime(or NumLeadsxTime depending on the resolution of the ECG)
+        return file
+
+
+

utils/ecg_utils.py

@@ -0,0 +1,114 @@
+# +
+import base64
+import struct
+
+import numpy as np
+import xmltodict
+from scipy.ndimage import median_filter as scipy_ndimage_median_filter
+import matplotlib.pyplot as plt
+import pywt
+# -
+
+lead_order = [
+    "I",
+    "II",
+    "III",
+    "aVR",
+    "aVL",
+    "aVF",
+    "V1",
+    "V2",
+    "V3",
+    "V4",
+    "V5",
+    "V6",
+]
+
+
+def plot_12_lead_ecg(ecg_array, output_filename=None):
+    """
+    Plot each lead of the 12-lead ECG, and save the plot to a file.
+    All leads share the x axis, but each lead gets its own chart.
+    """
+    fig, axs = plt.subplots(12, 1, sharex=True, figsize=(16, 9))
+    for lead, lead_name in enumerate(lead_order):
+        axs[lead].plot(ecg_array[:, lead])
+        axs[lead].set_ylabel(str(lead_name))
+    if output_filename is not None:
+        plt.savefig(output_filename)
+    plt.show()
+    plt.close()
+
+
+def get_median_filter_width(sampling_frequency, duration):
+    res = int(sampling_frequency * duration)
+    res += (res % 2) - 1  # needs to be an odd number
+    return res
+
+
+def remove_baseline_wander(waveform: np.ndarray, sampling_frequency: int) -> np.ndarray:
+
+    """
+    Remove baseline wander from ECG NPYs
+    de Chazal et al. used two median filters to remove baseline wander.
+    Median filters take the median value of a sliding window of a specified size
+    One median filter of 200-ms width to remove QRS complexes and P-waves and other of
+    600 ms width to remove T-waves.
+    Do one filter and then the next filter. Then take the result and subtract it form the original signal
+    https://pubmed.ncbi.nlm.nih.gov/15248536/
+    Example of median filter:
+    medfilt([2,6,5,4,0,3,5,7,9,2,0,1], 5) -> [ 2. 4. 4. 4. 4. 4. 5. 5. 5. 2. 1. 0.]
+    >>> np.median([0, 0, 2, 6, 5])
+    2.0
+    >>> np.median([0, 2, 6, 5, 4])
+    4.0
+
+    """
+
+    # Depending on the sampling frequency, the widths of the convolutional median filters changes
+    filter_widths = [
+        get_median_filter_width(sampling_frequency, duration) for duration in [0.2, 0.6]
+    ]
+    filter_widths = np.array(filter_widths, dtype="int")
+
+    # make a copy of orignal signal
+    original_waveform = waveform.copy()
+
+    # apply median filters one by one on top of each other
+    for filter_width in filter_widths:
+        waveform = scipy_ndimage_median_filter(
+            waveform, size=(filter_width, 1), mode="constant", cval=0.0
+        )
+    waveform = original_waveform - waveform  # finally subtract from orignal signal
+    return waveform
+
+
+def wavelet_denoise_signal(
+    waveform: np.ndarray,
+    dwt_transform: str = "bior4.4",
+    dlevels: int = 9,
+    cutoff_low: int = 1,
+    cutoff_high: int = 7,
+) -> np.ndarray:
+
+    # cutoff_low determines how flat you want overall baseline to be.
+    #   Higher means more flat baseline
+    # cutoff_high determines within the small segments how much do
+    #   you want to suppress the squiggliness. Lower cutoff_high
+    #   suppresses more squiggliness but also suppresses R wave morphology
+
+    coefficients = pywt.wavedec(
+        waveform, dwt_transform, level=dlevels
+    )  # wavelet transform 'bior4.4'
+    # scale 0 to cutoff_low
+    for low_cutoff_value in range(0, cutoff_low):
+        coefficients[low_cutoff_value] = np.multiply(
+            coefficients[low_cutoff_value], [0.0]
+        )
+    # scale cutoff_high to end
+    for high_cutoff_value in range(cutoff_high, len(coefficients)):
+        coefficients[high_cutoff_value] = np.multiply(
+            coefficients[high_cutoff_value], [0.0]
+        )
+    waveform = pywt.waverec(coefficients, dwt_transform)  # inverse wavelet transform
+    return waveform

utils/models.py

@@ -0,0 +1,247 @@
+from torch import nn
+import torch
+from functools import reduce
+from operator import __add__
+import torch.nn.functional as F
+from collections import OrderedDict
+from typing import Callable, List
+from torch import Tensor
+
+
+class EffNet(nn.Module):
+    # lightly retouched version of John's EffNet to add clean support for multiple output
+    # layer designs as well as single-lead inputs
+    def __init__(
+        self,
+        num_extra_inputs: int = 0,
+        output_neurons: int = 1,
+        channels: List[int] = (32, 16, 24, 40, 80, 112, 192, 320, 1280), 
+        depth: List[int] = (1, 2, 2, 3, 3, 3, 3),
+        dilation: int = 2,
+        stride: int = 8,
+        expansion: int = 6,
+        embedding_hook: bool = False,
+        input_channels: int = 1, 
+        verbose: bool = False,
+        embedding_shift: bool = False,
+    ):
+        super().__init__()
+
+        self.input_channels = input_channels
+        self.channels = channels
+        self.output_nerons = output_neurons
+
+        # backwards compatibility change to prevent the addition of the output_neurons param
+        # from breaking people's existing EffNet initializations
+        if len(self.channels) == 10:
+            self.output_nerons = self.channels[9]
+            print(
+                "DEPRECATION WARNING: instead of controlling the number of output neurons by changing the 10th item in the channels parameter, use the new output_neurons parameter instead."
+            )
+
+        self.depth = depth
+        self.expansion = expansion
+        self.stride = stride
+        self.dilation = dilation
+        self.embedding_hook = embedding_hook
+        self.embedding_shift = embedding_shift
+
+        if verbose:
+            print("\nEffNet Parameters:")
+            print(f"{self.input_channels=}")
+            print(f"{self.channels=}")
+            print(f"{self.output_nerons=}")
+            print(f"{self.depth=}")
+            print(f"{self.expansion=}")
+            print(f"{self.stride=}")
+            print(f"{self.dilation=}")
+            print(f"{self.embedding_hook=}")
+            print("\n")
+
+        self.stage1 = nn.Conv1d(
+            self.input_channels,
+            self.channels[0],
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            dilation=dilation,
+        )  # 1 conv
+
+        self.b0 = nn.BatchNorm1d(self.channels[0])
+
+        self.stage2 = MBConv(
+            self.channels[0], self.channels[1], self.expansion, self.depth[0], stride=2
+        )
+
+        self.stage3 = MBConv(
+            self.channels[1], self.channels[2], self.expansion, self.depth[1], stride=2
+        )
+
+        self.Pool = nn.MaxPool1d(3, stride=1, padding=1)
+
+        self.stage4 = MBConv(
+            self.channels[2], self.channels[3], self.expansion, self.depth[2], stride=2
+        )
+
+        self.stage5 = MBConv(
+            self.channels[3], self.channels[4], self.expansion, self.depth[3], stride=2
+        )
+
+        self.stage6 = MBConv(
+            self.channels[4], self.channels[5], self.expansion, self.depth[4], stride=2
+        )
+
+        self.stage7 = MBConv(
+            self.channels[5], self.channels[6], self.expansion, self.depth[5], stride=2
+        )
+
+        self.stage8 = MBConv(
+            self.channels[6], self.channels[7], self.expansion, self.depth[6], stride=2
+        )
+
+        self.stage9 = nn.Conv1d(self.channels[7], self.channels[8], kernel_size=1)
+        self.AAP = nn.AdaptiveAvgPool1d(1)
+        self.act = nn.ReLU()
+        self.drop = nn.Dropout(p=0.3)
+        self.num_extra_inputs = num_extra_inputs
+
+        self.fc = nn.Linear(self.channels[5] + num_extra_inputs, self.output_nerons)
+        self.fc = nn.Linear(self.channels[8] + num_extra_inputs, self.output_nerons)
+        
+        
+        self.fc.bias.data[0] = 0.275
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.num_extra_inputs > 0:
+            x, extra_inputs = x
+
+        x = self.b0(self.stage1(x))
+        x = self.stage2(x)
+        x = self.stage3(x)
+        x = self.Pool(x)
+        x = self.stage4(x)
+        x = self.stage5(x)
+        x = self.stage6(x)
+        x = self.Pool(x)
+        x = self.stage7(x)
+        x = self.stage8(x)
+        x = self.stage9(x)
+        x = self.act(self.AAP(x)[:, :, 0])
+        if self.embedding_hook:
+            return x
+        
+        else:
+        
+            if self.embedding_shift:
+                delta_embedding_array = np.load('/workspace/imin/applewatch_potassium/delta_embedding_poolaverage_5second_to_5second.npy')
+                delta_embedding_tensor = torch.tensor(delta_embedding_array, device='cuda')
+                x += delta_embedding_tensor
+
+            x = self.drop(x)
+
+            if self.num_extra_inputs > 0:
+                x = torch.cat((x, extra_inputs), 1)
+
+            x = self.fc(x)
+            return x
+
+
+class Bottleneck(nn.Module):
+    def __init__(
+        self,
+        in_channel: int,
+        out_channel: int,
+        expansion: int,
+        activation: Callable,
+        stride: int = 1,
+        padding: int = 1,
+    ):
+        super().__init__()
+
+        self.stride = stride
+        self.conv1 = nn.Conv1d(in_channel, in_channel * expansion, kernel_size=1)
+        self.conv2 = nn.Conv1d(
+            in_channel * expansion,
+            in_channel * expansion,
+            kernel_size=3,
+            groups=in_channel * expansion,
+            padding=padding,
+            stride=stride,
+        )
+        self.conv3 = nn.Conv1d(
+            in_channel * expansion, out_channel, kernel_size=1, stride=1
+        )
+        self.b0 = nn.BatchNorm1d(in_channel * expansion)
+        self.b1 = nn.BatchNorm1d(in_channel * expansion)
+        self.d = nn.Dropout()
+        self.act = activation()
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.stride == 1:
+            y = self.act(self.b0(self.conv1(x)))
+            y = self.act(self.b1(self.conv2(y)))
+            y = self.conv3(y)
+            y = self.d(y)
+            y = x + y
+            return y
+        else:
+            y = self.act(self.b0(self.conv1(x)))
+            y = self.act(self.b1(self.conv2(y)))
+            y = self.conv3(y)
+            return y
+
+
+class MBConv(nn.Module):
+    def __init__(
+        self, in_channel, out_channels, expansion, layers, activation=nn.ReLU6, stride=2
+    ):
+        super().__init__()
+
+        self.stack = OrderedDict()
+        for i in range(0, layers - 1):
+            self.stack["s" + str(i)] = Bottleneck(
+                in_channel, in_channel, expansion, activation
+            )
+
+        self.stack["s" + str(layers + 1)] = Bottleneck(
+            in_channel, out_channels, expansion, activation, stride=stride
+        )
+
+        self.stack = nn.Sequential(self.stack)
+
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.stack(x)
+        return self.bn(x)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels * BasicBlock.expansion, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels * BasicBlock.expansion)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = F.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = F.relu(out)
+
+        return out

utils/training_models.py

@@ -0,0 +1,447 @@
+from functools import lru_cache
+from typing import Iterable
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import numpy as np
+import wandb
+from sklearn import metrics as skl_metrics
+import torchvision
+import os
+from pathlib import Path
+import pandas as pd
+
+
+class TrainingMetric:
+    def __init__(self, metric_func, metric_name, optimum=None):
+        self.func = metric_func
+        self.name = metric_name
+        self.optimum = optimum
+
+    def calc_metric(self, *args, **kwargs):
+        try:
+            return self.func(*args, **kwargs)
+        except ValueError as e:
+            return np.nan
+
+    def __call__(self, y_true, y_pred, labels=None, split=None, step_type=None) -> dict:
+
+        # if y_true is empty
+        if y_true.shape[0] == 0:  # TODO: handle other cases
+            m = {
+                f"{step_type}_{split}_{l}_{self.name}": self.calc_metric(None, yp)
+                for yp, l in zip(y_pred.T, labels)
+            }
+            return m
+
+        # Simple 1:1 y_true and y_pred are either shape=(batch, 1) or shape=(batch,)
+        if len(y_pred.shape) == 1 or (y_pred.shape[1] == 1 and y_true.shape[1] == 1):
+            m = {
+                f"{step_type}_{split}_{self.name}": self.calc_metric(
+                    y_true.flatten(), y_pred.flatten()
+                )
+            }
+
+        # Multi-binary classification-like  y_true and y_pred are shape=(batch, class)
+        elif y_true.shape[1] != 1 and y_pred.shape[1] != 1:
+            m = {
+                f"{step_type}_{split}_{l}_{self.name}": self.calc_metric(yt, yp)
+                for yt, yp, l in zip(y_true.T, y_pred.T, labels)
+            }
+
+        # Multi-class classification-like  y_true is shape=(batch, 1) or shape=(batch,) and y_pred is shape=(batch, class)
+        elif (len(y_true.shape) == 1 or y_true.shape[1] == 1) and y_pred.shape[1] != 1:
+            m = {
+                f"{step_type}_{split}_{l}_{self.name}": self.calc_metric(
+                    y_true.flatten() == i, yp
+                )
+                for i, (yp, l) in enumerate(
+                    zip(y_pred.T, labels)
+                )  # turn multi class into binary classification
+            }
+
+        return m
+
+
+class CumulativeMetric(TrainingMetric):
+
+    """Wraps a metric to apply to every class in output and calculate a cumulative value (like mean AUC)"""
+
+    def __init__(
+        self,
+        training_metric: TrainingMetric,
+        metric_func,
+        metric_name="cumulative",
+        optimum=None,
+    ):
+        optimum = optimum or training_metric.optimum
+        metric_name = f"{metric_name}_{training_metric.name}"
+        super().__init__(metric_func, metric_name, optimum)
+        self.base_metric = training_metric
+
+    def __call__(self, y_true, y_pred, labels=None, split=None, step_type=None):
+        vals = list(self.base_metric(y_true, y_pred, labels, split, step_type).values())
+
+        m = {f"{step_type}_{split}_{self.name}": self.func(vals)}
+        return m
+
+
+r2_metric = TrainingMetric(skl_metrics.r2_score, "r2", optimum="max")
+roc_auc_metric = TrainingMetric(skl_metrics.roc_auc_score, "roc_auc", optimum="max")
+accuracy_metric = TrainingMetric(skl_metrics.accuracy_score, "accuracy", optimum="max")
+mae_metric = TrainingMetric(skl_metrics.mean_absolute_error, "mae", optimum="min")
+pred_value_mean_metric = TrainingMetric(
+    lambda y_true, y_pred: np.mean(y_pred), "pred_value_mean"
+)
+pred_value_std_metric = TrainingMetric(
+    lambda y_true, y_pred: np.std(y_pred), "pred_value_std"
+)
+
+
+class TrainingModel(pl.LightningModule):
+    def __init__(
+        self,
+        model,
+        metrics: Iterable[TrainingMetric] = dict(),
+        tracked_metric=None,
+        early_stop_epochs=10,
+        checkpoint_every_epoch=False,
+        checkpoint_every_n_steps=None,
+        index_labels=None,
+        save_predictions_path=None,
+        lr=0.01,
+    ):
+        super().__init__()
+        self.epoch_preds = {"train": ([], []), "val": ([], [])}
+        self.epoch_losses = {"train": [], "val": []}
+        self.metrics = {}
+        self.metric_funcs = {m.name: m for m in metrics}
+        self.tracked_metric = f"epoch_val_{tracked_metric}"
+        self.best_tracked_metric = None
+        self.early_stop_epochs = early_stop_epochs
+        self.checkpoint_every_epoch = checkpoint_every_epoch
+        self.checkpoint_every_n_steps = checkpoint_every_n_steps
+        self.metrics["epochs_since_last_best"] = 0
+        self.m = model
+        self.training_steps = 0
+        self.steps_since_checkpoint = 0
+        self.labels = index_labels
+        if self.labels is not None and isinstance(self.labels, str):
+            self.labels = [self.labels]
+        if isinstance(save_predictions_path, str):
+            save_predictions_path = Path(save_predictions_path)
+        self.save_predictions_path = save_predictions_path
+        self.lr = lr
+        self.step_loss = (None, None)
+
+        self.log_path = Path(wandb.run.dir) if wandb.run is not None else None
+
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.parameters(), self.lr)
+
+    def forward(self, x: dict):
+        # if anything other than 'primary_input' and 'extra_inputs' is used,
+        # this function must be overridden 
+        if 'extra_inputs' in x:
+            return self.m((x['primary_input'], x['extra_inputs']))
+        else:
+            return self.m(x['primary_input'])
+
+    def step(self, batch, step_type='train'):
+        batch = self.prepare_batch(batch)
+        y_pred = self.forward(batch)
+
+        if step_type != 'predict':
+            if 'labels' not in batch:
+                batch['labels'] = torch.empty(0)
+            loss = self.loss_func(y_pred, batch['labels'])
+            if torch.isnan(loss):
+                raise ValueError(loss)
+
+            self.log_step(step_type, batch['labels'], y_pred, loss)
+
+            return loss
+        else:
+            return y_pred
+    
+    def prepare_batch(self, batch):
+        return batch
+
+    def training_step(self, batch, i):
+        return self.step(batch, "train")
+
+    def validation_step(self, batch, i):
+        return self.step(batch, "val")
+
+    def predict_step(self, batch, *args):
+        y_pred = self.step(batch, "predict")
+        return {"filename": batch["filename"], "prediction": y_pred.cpu().numpy()}
+
+    def on_predict_epoch_end(self, results):
+
+        for i, predict_results in enumerate(results):
+            filename_df = pd.DataFrame(
+                {
+                    "filename": np.concatenate(
+                        [batch["filename"] for batch in predict_results]
+                    )
+                }
+            )
+
+            if self.labels is not None:
+                columns = [f"{class_name}_preds" for class_name in self.labels]
+            else:
+                columns = ["preds"]
+            outputs_df = pd.DataFrame(
+                np.concatenate(
+                    [batch["prediction"] for batch in predict_results], axis=0
+                ),
+                columns=columns,
+            )
+
+            prediction_df = pd.concat([filename_df, outputs_df], axis=1)
+
+            dataloader = self.trainer.predict_dataloaders[i]
+            manifest = dataloader.dataset.manifest
+            prediction_df = prediction_df.merge(manifest, on="filename", how="outer")
+            if wandb.run is not None:
+                prediction_df.to_csv(
+                    Path(wandb.run.dir).parent
+                    / "data"
+                    / f"dataloader_{i}_potassium_predictions.csv",
+                    index=False,
+                )
+            if self.save_predictions_path is not None:
+
+                if ".csv" in self.save_predictions_path.name:
+                    prediction_df.to_csv(
+                        self.save_predictions_path.parent
+                        / self.save_predictions_path.name.replace(".csv", f"_{i}_.csv"),
+                        index=False,
+                    )
+                else:
+                    prediction_df.to_csv(
+                        self.save_predictions_path / f"dataloader_{i}_potassium_predictions.csv",
+                        index=False,
+                    )
+
+            if wandb.run is None and self.save_predictions_path is None:
+                print(
+                    "WandB is not active and self.save_predictions_path is None. Predictions will be saved to the directory this script is being run in."
+                )
+                prediction_df.to_csv(f"dataloader_{i}_potassium_predictions.csv", index=False)
+
+    def log_step(self, step_type, labels, output_tensor, loss):
+        self.step_loss = (step_type, loss.detach().item())
+        self.epoch_preds[step_type][0].append(labels.detach().cpu().numpy())
+        self.epoch_preds[step_type][1].append(output_tensor.detach().cpu().numpy())
+        self.epoch_losses[step_type].append(loss.detach().item())
+        if step_type == "train":
+            self.training_steps += 1
+            self.steps_since_checkpoint += 1
+            if (
+                self.checkpoint_every_n_steps is not None
+                and self.steps_since_checkpoint > self.checkpoint_every_n_steps
+            ):
+                self.steps_since_checkpoint = 0
+                self.checkpoint_weights(f"step_{self.training_steps}")
+
+    def checkpoint_weights(self, name=""):
+        if wandb.run is not None:
+            weights_path = Path(wandb.run.dir).parent / "weights"
+            if not weights_path.is_dir():
+                weights_path.mkdir()
+            torch.save(self.state_dict(), weights_path / f"model_{name}.pt")
+        else:
+            print("Did not checkpoint model. wandb not initialized.")
+
+    def validation_epoch_end(self, preds):
+
+        # Save weights
+        self.metrics["epoch"] = self.current_epoch
+        if self.checkpoint_every_epoch:
+            self.checkpoint_weights(f"epoch_{self.current_epoch}")
+
+        # Calculate metrics
+        for m_type in ["train", "val"]:
+
+            y_true, y_pred = self.epoch_preds[m_type]
+            if len(y_true) == 0 or len(y_pred) == 0:
+                continue
+            y_true, y_pred = np.concatenate(y_true), np.concatenate(y_pred)
+
+            self.metrics[f"epoch_{m_type}_loss"] = np.mean(self.epoch_losses[m_type])
+            for m in self.metric_funcs.values():
+                self.metrics.update(
+                    m(
+                        y_true,
+                        y_pred,
+                        labels=self.labels,
+                        split=m_type,
+                        step_type="epoch",
+                    )
+                )
+
+            # Reset predictions
+            self.epoch_losses[m_type] = []
+            self.epoch_preds[m_type] = ([], [])
+
+        # Check if new best epoch
+        if self.metrics is not None and self.tracked_metric is not None:
+            if self.tracked_metric == "epoch_val_loss":
+                metric_optimization = "min"
+            else:
+                metric_optimization = self.metric_funcs[
+                    self.tracked_metric.replace("epoch_val_", "")
+                ].optimum
+            if (
+                self.metrics[self.tracked_metric] is not None
+                and (
+                    self.best_tracked_metric is None
+                    or (
+                        metric_optimization == "max"
+                        and self.metrics[self.tracked_metric] > self.best_tracked_metric
+                    )
+                    or (
+                        metric_optimization == "min"
+                        and self.metrics[self.tracked_metric] < self.best_tracked_metric
+                    )
+                )
+                and self.current_epoch > 0
+            ):
+                print(
+                    f"New best epoch! {self.tracked_metric}={self.metrics[self.tracked_metric]}, epoch={self.current_epoch}"
+                )
+                self.checkpoint_weights(f"best_{self.tracked_metric}")
+                self.metrics["epochs_since_last_best"] = 0
+                self.best_tracked_metric = self.metrics[self.tracked_metric]
+            else:
+                self.metrics["epochs_since_last_best"] += 1
+            if self.metrics["epochs_since_last_best"] >= self.early_stop_epochs:
+                raise KeyboardInterrupt("Early stopping condition met")
+
+        # Log to w&b
+        if wandb.run is not None:
+            wandb.log(self.metrics)
+
+
+class RegressionModel(TrainingModel):
+    def __init__(
+        self,
+        model,
+        metrics=(r2_metric, mae_metric, pred_value_mean_metric, pred_value_std_metric),
+        tracked_metric="mae",
+        early_stop_epochs=10,
+        checkpoint_every_epoch=False,
+        checkpoint_every_n_steps=None,
+        index_labels=None,
+        save_predictions_path=None,
+        lr=0.01,
+    ):
+        super().__init__(
+            model=model,
+            metrics=metrics,
+            tracked_metric=tracked_metric,
+            early_stop_epochs=early_stop_epochs,
+            checkpoint_every_epoch=checkpoint_every_epoch,
+            checkpoint_every_n_steps=checkpoint_every_n_steps,
+            index_labels=index_labels,
+            save_predictions_path=save_predictions_path,
+            lr=lr,
+        )
+        self.loss_func = nn.MSELoss()
+
+    def prepare_batch(self, batch):
+        if "labels" in batch and len(batch["labels"].shape) == 1:
+            batch["labels"] = batch["labels"][:, None]
+        return batch
+
+
+class BinaryClassificationModel(TrainingModel):
+    def __init__(
+        self,
+        model,
+        metrics=(roc_auc_metric, CumulativeMetric(roc_auc_metric, np.nanmean, "mean")),
+        tracked_metric="mean_roc_auc",
+        early_stop_epochs=10,
+        checkpoint_every_epoch=False,
+        checkpoint_every_n_steps=None,
+        index_labels=None,
+        save_predictions_path=None,
+        lr=0.01,
+    ):
+        super().__init__(
+            model=model,
+            metrics=metrics,
+            tracked_metric=tracked_metric,
+            early_stop_epochs=early_stop_epochs,
+            checkpoint_every_epoch=checkpoint_every_epoch,
+            checkpoint_every_n_steps=checkpoint_every_n_steps,
+            index_labels=index_labels,
+            save_predictions_path=save_predictions_path,
+            lr=lr,
+        )
+        self.loss_func = nn.BCEWithLogitsLoss()
+
+    def prepare_batch(self, batch):
+        if "labels" in batch and len(batch["labels"].shape) == 1:
+            batch["labels"] = batch["labels"][:, None]
+        return batch
+
+
+# Addresses bug caused by labels from a single column in a manifest being delivered as Bx1,
+# but nn.CrossEntropyLoss wants a simple list of length B.
+class SqueezeCrossEntropyLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.cross_entropy = nn.CrossEntropyLoss()
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor):
+        return self.cross_entropy(y_pred, y_true.squeeze(dim=-1))
+
+
+class MultiClassificationModel(TrainingModel):
+    def __init__(
+        self,
+        model,
+        metrics=(roc_auc_metric, CumulativeMetric(roc_auc_metric, np.mean, "mean")),
+        tracked_metric="mean_roc_auc",
+        early_stop_epochs=10,
+        checkpoint_every_epoch=False,
+        checkpoint_every_n_steps=None,
+        index_labels=None,
+        save_predictions_path=None,
+        lr=0.01,
+    ):
+        metrics = [*metrics]
+        super().__init__(
+            model=model,
+            metrics=metrics,
+            tracked_metric=tracked_metric,
+            early_stop_epochs=early_stop_epochs,
+            checkpoint_every_epoch=checkpoint_every_epoch,
+            checkpoint_every_n_steps=checkpoint_every_n_steps,
+            index_labels=index_labels,
+            save_predictions_path=save_predictions_path,
+            lr=lr,
+        )
+        self.loss_func = SqueezeCrossEntropyLoss()
+
+    def prepare_batch(self, batch):
+        if "labels" in batch:
+            batch["labels"] = batch["labels"].long()
+        batch["primary_input"] = batch["primary_input"].float()
+        return batch
+
+
+if __name__ == "__main__":
+    os.environ["WANDB_MODE"] = "offline"
+
+    m = torchvision.models.video.r2plus1d_18()
+    m.fc = nn.Linear(512, 1)
+    training_model = RegressionModel(m)
+    x = torch.randn((4, 3, 8, 112, 112))
+    y = m(x)
+    print(y.shape)
+
+