Add feature precision variation transform

Author: LocalToasty

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stamp"
-version = "2.0.0-dev6"
+version = "2.0.0-dev7"
 authors = [
     { name = "Omar El Nahhas", email = "omar.el_nahhas@tu-dresden.de" },
     { name = "Marko van Treeck", email = "markovantreeck@gmail.com" },

src/stamp/__main__.py

@@ -106,6 +106,7 @@ def _run_cli(args: argparse.Namespace) -> None:
                 max_epochs=config.training.max_epochs,
                 patience=config.training.patience,
                 accelerator=config.training.accelerator,
+                use_vary_precision_transform=config.training.use_vary_precision_transform,
             )
 
         case "deploy":
@@ -156,11 +157,13 @@ def _run_cli(args: argparse.Namespace) -> None:
                 # Dataset and -loader parameters
                 bag_size=config.crossval.bag_size,
                 num_workers=config.crossval.num_workers,
-                # crossval paramenters
+                # Crossval paramenters
                 batch_size=config.crossval.batch_size,
                 max_epochs=config.crossval.max_epochs,
                 patience=config.crossval.patience,
                 accelerator=config.crossval.accelerator,
+                # Experimental Features
+                use_vary_precision_transform=config.crossval.use_vary_precision_transform,
             )
 
         case "statistics":

src/stamp/config.yaml

@@ -104,6 +104,14 @@ training:
   # If unspecified, they will be inferred from the table itself.
   #categories: ["mutated", "wild type"]
 
+  # Experimental features:
+
+  # Please try uncommenting the settings below
+  # and report if they improve / reduce model performance!
+
+  # Change the precision of features during training
+  #use_vary_precision_transform: true
+
 
 deployment:
   output_dir: "/path/to/save/files/to"

src/stamp/modeling/config.py

@@ -34,6 +34,9 @@ class TrainConfig(BaseModel):
     patience: int = 16
     accelerator: str = "gpu" if torch.cuda.is_available() else "cpu"
 
+    # Experimental features
+    use_vary_precision_transform: bool = False
+
 
 class CrossvalConfig(TrainConfig):
     n_splits: int = Field(5, ge=2)

src/stamp/modeling/crossval.py

@@ -22,6 +22,7 @@
 from stamp.modeling.deploy import _predict, _to_prediction_df
 from stamp.modeling.lightning_model import LitVisionTransformer
 from stamp.modeling.train import setup_model_for_training, train_model_
+from stamp.modeling.transforms import VaryPrecisionTransform
 
 __author__ = "Marko van Treeck"
 __copyright__ = "Copyright (C) 2024 Marko van Treeck"
@@ -57,6 +58,8 @@ def categorical_crossval_(
     max_epochs: int,
     patience: int,
     accelerator: str | Accelerator,
+    # Experimental features
+    use_vary_precision_transform: bool,
 ) -> None:
     patient_to_ground_truth: Final[dict[PatientId, GroundTruth]] = (
         patient_to_ground_truth_from_clini_table_(
@@ -149,6 +152,11 @@ def categorical_crossval_(
                         }
                     )
                 ),
+                train_transform=(
+                    VaryPrecisionTransform(min_fraction_bits=1)
+                    if use_vary_precision_transform
+                    else None
+                ),
             )
             model = train_model_(
                 output_dir=split_dir,
@@ -203,4 +211,4 @@ def _get_splits(
             )
         ]
     )
-    return splits
+    return splits

src/stamp/modeling/data.py

@@ -1,7 +1,7 @@
 """Helper classes to manage pytorch data."""
 
 import logging
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import KW_ONLY, dataclass
 from itertools import groupby
 from pathlib import Path
@@ -63,6 +63,7 @@ def dataloader_from_patient_data(
     batch_size: int,
     shuffle: bool,
     num_workers: int,
+    transform: Callable[[Tensor], Tensor] | None,
 ) -> tuple[DataLoader[tuple[Bags, BagSizes, EncodedTargets]], Sequence[Category]]:
     """Creates a dataloader from patient data, encoding the ground truths.
 
@@ -81,6 +82,7 @@ def dataloader_from_patient_data(
         bags=[patient.feature_files for patient in patient_data],
         bag_size=bag_size,
         ground_truths=one_hot,
+        transform=transform,
     )
 
     return (
@@ -133,6 +135,8 @@ class BagDataset(Dataset[tuple[_Bag, BagSize, _EncodedTarget]]):
     ground_truths: Bool[Tensor, "index category_is_hot"]
     """The ground truth for each bag, one-hot encoded."""
 
+    transform: Callable[[Tensor], Tensor] | None
+
     def __post_init__(self) -> None:
         if len(self.bags) != len(self.ground_truths):
             raise ValueError(
@@ -152,8 +156,11 @@ def __getitem__(self, index: int) -> tuple[_Bag, BagSize, _EncodedTarget]:
                 )
         feats = torch.concat(feats).float()
 
+        if self.transform is not None:
+            feats = self.transform(feats)
+
         # Sample a subset, if required
-        if self.bag_size:
+        if self.bag_size is not None:
             return (
                 *_to_fixed_size_bag(feats, bag_size=self.bag_size),
                 self.ground_truths[index],
@@ -166,7 +173,7 @@ def __getitem__(self, index: int) -> tuple[_Bag, BagSize, _EncodedTarget]:
             )
 
 
-def _to_fixed_size_bag(bag: _Bag, bag_size: BagSize = 512) -> tuple[_Bag, BagSize]:
+def _to_fixed_size_bag(bag: _Bag, bag_size: BagSize) -> tuple[_Bag, BagSize]:
     """Samples a fixed-size bag of tiles from an arbitrary one.
 
     If the original bag did not have enough tiles,

src/stamp/modeling/deploy.py

@@ -126,6 +126,7 @@ def _predict(
         batch_size=1,
         shuffle=False,
         num_workers=num_workers,
+        transform=None,
     )
 
     trainer = lightning.Trainer(
@@ -136,9 +137,7 @@ def _predict(
     predictions = torch.concat(
         cast(
             list[torch.Tensor],
-            trainer.predict(
-                cast(lightning.LightningModule, torch.compile(model)), test_dl
-            ),
+            trainer.predict(model, test_dl),
         )
     )
 

src/stamp/modeling/train.py

@@ -1,5 +1,5 @@
 import shutil
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 from typing import cast
 
@@ -32,6 +32,7 @@
     EncodedTargets,
     LitVisionTransformer,
 )
+from stamp.modeling.transforms import VaryPrecisionTransform
 
 __author__ = "Marko van Treeck"
 __copyright__ = "Copyright (C) 2024 Marko van Treeck"
@@ -56,6 +57,8 @@ def train_categorical_model_(
     max_epochs: int,
     patience: int,
     accelerator: str | Accelerator,
+    # Experimental features
+    use_vary_precision_transform: bool,
 ) -> None:
     """Trains a model.
 
@@ -119,6 +122,11 @@ def train_categorical_model_(
         clini_table=clini_table,
         slide_table=slide_table,
         feature_dir=feature_dir,
+        train_transform=(
+            VaryPrecisionTransform(min_fraction_bits=1)
+            if use_vary_precision_transform
+            else None
+        ),
     )
     train_model_(
         output_dir=output_dir,
@@ -187,6 +195,7 @@ def setup_model_for_training(
     bag_size: int,
     batch_size: int,
     num_workers: int,
+    train_transform: Callable[[torch.Tensor], torch.Tensor] | None,
     # Metadata, has no effect on model training
     ground_truth_label: PandasLabel,
     clini_table: Path,
@@ -225,6 +234,7 @@ def setup_model_for_training(
         batch_size=batch_size,
         shuffle=True,
         num_workers=num_workers,
+        transform=train_transform,
     )
     del categories  # Let's not accidentally reuse the original categories
     valid_dl, _ = dataloader_from_patient_data(
@@ -234,6 +244,7 @@ def setup_model_for_training(
         batch_size=1,
         shuffle=False,
         num_workers=num_workers,
+        transform=None,
     )
     if overlap := set(train_patients) & set(valid_patients):
         raise RuntimeError(

src/stamp/modeling/transforms.py

@@ -0,0 +1,41 @@
+import torch
+from jaxtyping import Float
+
+
+def vary_precision(
+    data: Float[torch.Tensor, "*dims"], *, min_fraction_bits: int
+) -> Float[torch.Tensor, "*dims"]:
+    """Randomly reduces the precision of the tensor's values."""
+    if min_fraction_bits < 1:
+        raise ValueError("min_fraction bits has to be at least 1")
+
+    if data.dtype == torch.float32:
+        fraction_bits = 23
+        mask_dtype = torch.int32
+    elif data.dtype == torch.float16:
+        fraction_bits = 10
+        mask_dtype = torch.int16
+    elif data.dtype == torch.bfloat16:
+        fraction_bits = 7
+        mask_dtype = torch.int16
+    else:
+        raise NotImplementedError(
+            f"precision variation not implemented for {data.dtype}"
+        )
+
+    no_of_bits_to_mask = torch.randint(0, fraction_bits - min_fraction_bits, data.shape)
+    mask = (~0 << no_of_bits_to_mask).to(dtype=mask_dtype, device=data.device)
+    augmented = (data.view(mask_dtype) & mask).view(data.dtype)
+    return augmented
+
+
+class VaryPrecisionTransform:
+    """A transform randomly reducing the precision of its inputs."""
+
+    def __init__(self, *, min_fraction_bits: int = 1) -> None:
+        self.min_fraction_bits = min_fraction_bits
+
+    def __call__(
+        self, batch: Float[torch.Tensor, "*dims"]
+    ) -> Float[torch.Tensor, "*dims"]:
+        return vary_precision(data=batch, min_fraction_bits=self.min_fraction_bits)

tests/test_data.py

@@ -76,7 +76,6 @@ def test_dataset(
     dim_feats: int = 34,
     batch_size: int = 2,
 ) -> None:
-
     ds = BagDataset(
         bags=[
             [_make_feature_file(torch.rand((12, dim_feats)))],
@@ -85,6 +84,7 @@ def test_dataset(
         ],
         bag_size=bag_size,
         ground_truths=torch.rand(3, 4) > 0.5,
+        transform=None,
     )
 
     assert len(ds) == 3