Size tensor embeddings without re-processing samples

docs/api/processors/pyhealth.processors.TensorProcessor.rst

@@ -3,6 +3,23 @@ pyhealth.processors.TensorProcessor
 
 Processor for tensor data.
 
+``size()`` returns the width of the feature's last dimension seen in ``fit()``
+(1 for scalars), or ``None`` before fitting.
+
+Custom tensor processors
+------------------------
+
+Subclass ``TensorProcessor`` to transform numeric features, for example to select
+columns, impute or scale. Models such as ``MLP`` size their input layer from the
+already-processed samples, falling back to ``size()``, and never call
+``process()`` on a sample a second time, so ``process()`` does not need to be
+idempotent. If ``process()`` changes the width, override ``size()`` to return the
+output width. See ``examples/custom_tensor_processor.py``.
+
+In PyHealth 2.0.2 and earlier, ``EmbeddingModel`` called ``process()`` again on
+already-processed samples to infer the width, which broke processors that are not
+idempotent.
+
 .. autoclass:: pyhealth.processors.TensorProcessor
     :members:
     :undoc-members:

examples/custom_tensor_processor.py

@@ -0,0 +1,63 @@
+"""A custom tensor processor that changes the feature width, used with an MLP.
+
+Subclass ``TensorProcessor`` to transform numeric features, here keeping a
+subset of columns. Models size their input layer from the already-processed
+samples (or from ``size()``), and never call ``process()`` a second time, so a
+processor like this one, which is not idempotent, works with any model.
+
+Runs in a few seconds on CPU with synthetic data; no download.
+
+Usage:
+    python examples/custom_tensor_processor.py
+"""
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset, get_dataloader
+from pyhealth.models import MLP
+from pyhealth.processors import TensorProcessor
+
+
+class SelectColumns(TensorProcessor):
+    """Keeps the given columns of a numeric feature vector."""
+
+    def __init__(self, columns):
+        super().__init__()
+        self.columns = list(columns)
+
+    def process(self, value):
+        return super().process(value)[..., self.columns]
+
+    def size(self):
+        return len(self.columns)
+
+
+def main():
+    torch.manual_seed(0)
+    samples = [
+        {
+            "patient_id": f"p{i}",
+            # age, heart rate, temperature, a column we want to drop
+            "vitals": [40.0 + i, 70.0 + i % 9, 36.5 + (i % 3) / 10, 999.0],
+            "label": int(i % 3 == 0),
+        }
+        for i in range(30)
+    ]
+    dataset = create_sample_dataset(
+        samples=samples,
+        input_schema={"vitals": SelectColumns(columns=[0, 1, 2])},
+        output_schema={"label": "binary"},
+        dataset_name="custom_tensor_processor",
+    )
+    print("processed sample:", dataset[0]["vitals"])
+
+    model = MLP(dataset=dataset)
+    print("input layer:", model.embedding_model.embedding_layers["vitals"])
+
+    batch = next(iter(get_dataloader(dataset, batch_size=8)))
+    out = model(**batch)
+    print("y_prob shape:", tuple(out["y_prob"].shape), "loss:", round(float(out["loss"]), 4))
+
+
+if __name__ == "__main__":
+    main()

pyhealth/models/embedding.py

@@ -133,9 +133,26 @@ class EmbeddingModel(BaseModel):
         Input: (..., size)
         Output: (..., embedding_dim)
 
-    - TensorProcessor: nn.Linear (size inferred from first sample)
+    - TensorProcessor: nn.Linear, sized from the first already-processed
+        sample (or the processor's ``size()``); ``process()`` is not re-run
 
     - MultiHotProcessor: nn.Linear over multi-hot vector
+
+    Examples:
+        >>> from pyhealth.datasets import create_sample_dataset
+        >>> from pyhealth.models import EmbeddingModel
+        >>> samples = [
+        ...     {"patient_id": "p0", "x": [1.0, 2.0, 3.0], "label": 0},
+        ...     {"patient_id": "p1", "x": [4.0, 5.0, 6.0], "label": 1},
+        ... ]
+        >>> dataset = create_sample_dataset(
+        ...     samples=samples,
+        ...     input_schema={"x": "tensor"},
+        ...     output_schema={"label": "binary"},
+        ... )
+        >>> model = EmbeddingModel(dataset, embedding_dim=8)
+        >>> model.embedding_layers["x"]
+        Linear(in_features=3, out_features=8, bias=True)
     """
 
     def __init__(
@@ -213,16 +230,18 @@ def __init__(
                 )
 
             elif isinstance(processor, TensorProcessor):
-                # Infer size from first sample
-                sample_tensor = None
+                # Dataset samples are already processed: read the width from the
+                # first one. Never call process() again, since processors that
+                # select columns, scale or impute are not idempotent.
+                input_size = None
                 for sample in dataset:
                     if field_name in sample:
-                        sample_tensor = processor.process(sample[field_name])
+                        value = torch.as_tensor(sample[field_name])
+                        input_size = value.shape[-1] if value.dim() > 0 else 1
                         break
-                if sample_tensor is not None:
-                    input_size = (
-                        sample_tensor.shape[-1] if sample_tensor.dim() > 0 else 1
-                    )
+                if input_size is None:
+                    input_size = processor.size()
+                if input_size is not None:
                     self.embedding_layers[field_name] = nn.Linear(
                         in_features=input_size, out_features=embedding_dim
                     )

pyhealth/processors/tensor_processor.py

@@ -19,6 +19,15 @@ class TensorProcessor(FeatureProcessor):
 
     Output:
         - torch.Tensor with appropriate shape and dtype
+
+    Examples:
+        >>> from pyhealth.processors import TensorProcessor
+        >>> processor = TensorProcessor()
+        >>> processor.fit([{"x": [1.0, 2.0, 3.0]}], "x")
+        >>> processor.size()
+        3
+        >>> processor.process([4.0, 5.0, 6.0])
+        tensor([4., 5., 6.])
     """
 
     def __init__(
@@ -37,10 +46,11 @@ def __init__(
         """
         self.dtype = dtype
         self._n_dim = None
+        self._feature_dim = None
         self._spatial_dims = spatial_dims
 
     def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None:
-        """Infer n_dim from the first valid sample.
+        """Infer n_dim and the feature width from the first valid sample.
 
         Args:
             samples: Iterable of sample dictionaries.
@@ -55,6 +65,7 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None:
                     else torch.tensor(value, dtype=self.dtype)
                 )
                 self._n_dim = tensor.dim()
+                self._feature_dim = tensor.shape[-1] if tensor.dim() > 0 else 1
                 break
 
     def process(self, value: Any) -> torch.Tensor:
@@ -74,14 +85,15 @@ def process(self, value: Any) -> torch.Tensor:
             return value.detach().clone().to(dtype=self.dtype)
         return torch.tensor(value, dtype=self.dtype)
 
-    def size(self) -> None:
+    def size(self) -> int | None:
         """
         Get the feature size of the processor.
 
         Returns:
-            None: Size is not predetermined for tensor processor
+            The width of the last dimension (1 for scalars) seen in ``fit()``,
+            or None before fitting.
         """
-        return None
+        return getattr(self, "_feature_dim", None)
 
     def is_token(self) -> bool:
         """Whether the output tensor represents discrete token indices, inferred from dtype.

tests/core/test_embedding_tensor_width.py

@@ -0,0 +1,80 @@
+"""EmbeddingModel must size tensor features without re-processing samples.
+
+Samples read from a SampleDataset are already processed. Calling
+``processor.process()`` on them again is wrong for any processor that is not
+idempotent (one that selects columns, scales or imputes).
+"""
+
+import unittest
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset
+from pyhealth.models import MLP
+from pyhealth.processors import TensorProcessor
+
+
+def _samples():
+    return [
+        {"patient_id": f"p{i}", "x": [float(i), 1.0, 2.0, 3.0], "label": i % 2}
+        for i in range(6)
+    ]
+
+
+class FirstTwoColumns(TensorProcessor):
+    """Keeps the first two columns: output width 2, not idempotent."""
+
+    def process(self, value):
+        tensor = super().process(value)
+        if tensor.shape[-1] != 4:
+            raise ValueError("process() called on an already-processed sample")
+        return tensor[..., :2]
+
+    def size(self):
+        return 2
+
+
+class TestTensorProcessorSize(unittest.TestCase):
+    def test_fit_records_feature_width(self):
+        p = TensorProcessor()
+        p.fit(_samples(), "x")
+        self.assertEqual(p.size(), 4)
+
+    def test_size_is_none_before_fit(self):
+        self.assertIsNone(TensorProcessor().size())
+
+    def test_scalar_feature_has_width_one(self):
+        p = TensorProcessor()
+        p.fit([{"x": 3.0}], "x")
+        self.assertEqual(p.size(), 1)
+
+
+class TestEmbeddingTensorWidth(unittest.TestCase):
+    def test_builds_without_reprocessing_samples(self):
+        dataset = create_sample_dataset(
+            samples=_samples(),
+            input_schema={"x": FirstTwoColumns()},
+            output_schema={"label": "binary"},
+            dataset_name="test_embedding_tensor_width",
+        )
+        self.assertEqual(tuple(dataset[0]["x"].shape), (2,))
+        model = MLP(dataset=dataset)  # raised ValueError before the fix
+        layer = model.embedding_model.embedding_layers["x"]
+        self.assertEqual(layer.in_features, 2)
+        batch = {"x": torch.stack([dataset[i]["x"] for i in range(3)])}
+        out = model.embedding_model(batch)
+        self.assertEqual(out["x"].shape[0], 3)
+
+    def test_plain_tensor_feature_width(self):
+        dataset = create_sample_dataset(
+            samples=_samples(),
+            input_schema={"x": "tensor"},
+            output_schema={"label": "binary"},
+            dataset_name="test_embedding_tensor_width_plain",
+        )
+        model = MLP(dataset=dataset)
+        self.assertEqual(model.embedding_model.embedding_layers["x"].in_features, 4)
+
+
+if __name__ == "__main__":
+    unittest.main()